diff --git a/docs/commands/build.md b/docs/commands/build.md index 56e9434ab..9453f1f87 100644 --- a/docs/commands/build.md +++ b/docs/commands/build.md @@ -22,6 +22,7 @@ $ winml build [options] |---|---|---|---|---| | `--config` | `-c` | path | `None` | `WinMLBuildConfig` JSON file, generated by `winml config`. If omitted, config is auto-generated from `-m`. | | `--model` | `-m` | string | `None` | Hugging Face model ID or path to an existing `.onnx` file. | +| `--backend` | | choice | `None` | Backend for auto-generated config. `cgc` selects CGC preparation and CGIR conversion; it cannot be combined with `--ep`. | | `--export-type` | | choice | `generic` | Output selector: `generic` builds the stock single/composite ONNX model; `optimized` builds the family's registered runtime-optimized recipe (today the onnxruntime-genai CPU/NPU bundle) for the **resolved** `--ep`/`--device`. `optimized` fails fast if the family has no recipe or the resolved target is not one the recipe supports. | | `--output-dir` | `-o` | path | `None` | Directory for all build artifacts. Mutually exclusive with `--use-cache`. | | `--use-cache/--no-use-cache` | | flag | `false` | Store artifacts in the winml-cli global cache (`~/.cache/winml/`). Mutually exclusive with `--output-dir`. | @@ -57,6 +58,38 @@ single-pass build. Individual stages can be suppressed with `--no-quant`, !!! tip "Reproducible CI/CD builds" The config file is a portable, self-contained pipeline specification. Check it into source control and invoke `winml build -c config.json` in CI to produce identical artifacts without manual flag management. Set `"auto": false` in the config to disable the autoconf discovery loop for fully deterministic output. +## Optional ONNX-to-MLIR conversion + +Use `winml build -m model.onnx --backend cgc -o output` to auto-generate +the CGC configuration and build MLIR directly. Use `--ep winmlcg` instead +to prepare ONNX for WinMLCG EP without converting it to MLIR. These two +target selectors cannot be combined. + +Add a `convert` section to the build config to export the final ONNX model as CGIR: + +```json +{ + "compile": null, + "convert": { + "target": "cgir", + "options": { + "external_weights": true + } + } +} +``` + +This fragment supplements the existing build config. Export, optimization, +quantization, and compilation retain their existing configuration controls; +conversion does not disable compilation automatically. Set `compile` to `null` +when compilation is not wanted. + +The CLI retains `model.onnx` and writes `model.mlir` with the exporter sidecars. +Conversion also runs when the ONNX build is reused. Omit `convert` or set it to +`null` to preserve the ordinary ONNX build. This stage runs for single-model and +composite CLI builds. Module-mode arrays and optimized GenAI bundles accept the +configuration but do not execute the conversion stage. + ## Genai bundles for decoder LLMs (CPU/NPU) For a registered decoder-LLM family (currently **Qwen3**), `--export-type diff --git a/docs/commands/config.md b/docs/commands/config.md index 278fe3220..7cda9905c 100644 --- a/docs/commands/config.md +++ b/docs/commands/config.md @@ -29,6 +29,7 @@ $ winml config [options] | `--device` | `-d` | `auto\|npu\|gpu\|cpu` | `auto` | Target device. Affects the generated quantization and compilation sub-configs. `auto` leaves those sections unchanged from the kit defaults. | | `--ep` | | `TEXT` | *(none)* | Force a specific execution provider (`qnn`, `dml`, `migraphx`, `tensorrt`, `vitisai`, `openvino`, `cpu`). Overrides the device-to-provider mapping. When used without `--device`, the device is inferred from the EP. | | `--precision` | `-p` | `TEXT` | `auto` | Target precision: `auto`, `fp32`, `fp16`, `int8`, `int16`, or a mixed format such as `w8a16`. `auto` selects the precision based on the chosen device. | +| `--backend` | | `ort\|cgc` | *(none)* | `cgc` enables CGC compatibility rules, default FP16 conversion, no compilation, and a CGIR convert stage. Cannot combine `cgc` with `--ep`. Omission or `ort` preserves existing behavior. | | `--output` | `-o` | `PATH` | *(stdout)* | Write the generated JSON to this file instead of printing to stdout. | | `--library` | | `TEXT` | `transformers` | Source library for `TasksManager` task lookup. Defaults to `transformers`; set to `diffusers` or another Optimum-supported library when needed. | | `--quant/--no-quant` | | flag | `true` | Include quantization in the generated config (use `--no-quant` to omit it and set `quant` to `null`). | @@ -39,6 +40,37 @@ $ winml config [options] `winml config` queries the HuggingFace `TasksManager` to auto-detect the model's task, class, and ONNX export specification. For known model types it looks up a per-model kit in `MODEL_BUILD_CONFIGS` and uses that as a starting point, layering in your device, precision, and override file on top. When `-m` points to an existing `.onnx` file, the export stage is skipped by setting `export` to `null` in the output. The result is a complete `WinMLBuildConfig` JSON printed to stdout or written to a file, ready to be passed to `winml build`. +## CGC configuration + +The config generators use CGC stage settings for `--backend cgc` or +`--ep winmlcg`: `auto: false`, all registered CGC compatibility rules in +`optim`, `optim.ort_graph_optimization: false` to skip ORT graph optimization, +default `quant.mode: "fp16"`, and `compile: null`. +Only `--backend cgc` automatically adds `convert.target: "cgir"`. +Choose either `--backend cgc` or `--ep winmlcg`; combining `--backend cgc` +with any `--ep` is a usage error. +Explicit `--precision` values are preserved instead of forcing FP16. +Existing QDQ ONNX inputs and `--no-quant` use `quant: null`. +Loader and PyTorch-to-ONNX export settings are unchanged; ONNX input retains +`export: null`. + +The compatibility preset includes opset deduplication, scalar initializer Cast +folding, empty Resize input omission, Tile repeats materialization, supported +Resize coordinate conversion, cubic-to-linear Resize approximation, identity +GatherND reshaping, PRelu decomposition, and static DFT decomposition. +**Cubic-to-linear Resize is lossy and may reduce accuracy.** The optimizer logs +a warning when it applies this rule. Review accuracy after building; disable +`optim.approximate_cubic_resize_with_linear` in the generated JSON when needed. +No `--runtime` option is needed for config generation: Runtime CGC and the +WinMLCG EP share this offline compatibility and FP16 preparation. + +```bash +winml config -m microsoft/resnet-50 --backend cgc -o config.json +``` + +Module and composite configurations receive the same settings. Module-mode build +currently accepts but does not execute the convert stage. + ## Examples Generate a config for ResNet-50 with all auto-detected settings: diff --git a/docs/commands/eval.md b/docs/commands/eval.md index 164ff1fe2..8cb9faab0 100644 --- a/docs/commands/eval.md +++ b/docs/commands/eval.md @@ -1,6 +1,6 @@ # winml eval -> Evaluate ONNX or native Hugging Face PyTorch model accuracy on a standard dataset. +> Evaluate ONNX, CGC MLIR, or native Hugging Face PyTorch models. ## When to use this @@ -16,8 +16,8 @@ $ winml eval [options] | Flag | Short | Type | Default | Description | |---|---|---|---|---| -| `--model` | `-m` | `TEXT` | — | HuggingFace model ID, or path to a local `.onnx` file. Required (unless `--model-id` is provided directly). | -| `--model-id` | | `TEXT` | — | HuggingFace model ID used for preprocessor and config resolution when `-m` points to an `.onnx` file. Required when `-m` is an ONNX file. | +| `--model` | `-m` | `TEXT` | — | HuggingFace model ID, or path to a local `.onnx` or `.mlir` file. Required (unless `--model-id` is provided directly). | +| `--model-id` | | `TEXT` | — | HuggingFace model ID used for preprocessor and config resolution when `-m` points to an ONNX or MLIR file. | | `--task` | | `TEXT` | auto-detected | Task name (e.g., `image-classification`). Auto-detected from `--model-id` when not provided. Required when `-m` is an ONNX file and the task cannot be inferred. | | `--precision` | | `TEXT` | `auto` | Precision used when building the model from a HuggingFace ID. One of `auto`, `fp32`, `fp16`, `int8`, `int16`, or a mixed `w{x}a{y}` spec (e.g., `w8a16`). `fp16`/`fp32` skip quantization. **Ignored** when `-m` is a pre-built `.onnx` file — the precision is already baked in. | | `--device` | | choice | `auto` | Target device. Choices: `auto`, `npu`, `gpu`, `cpu`. `auto` selects the best available device. Combined with `--precision`, this drives the build when `-m` is a HuggingFace ID. | @@ -27,7 +27,7 @@ $ winml eval [options] | `--input-specs` | | `PATH` | — | JSON input tensor specs to merge into the Hugging Face export config. Symbolic string dimensions infer dynamic axes. **Ignored for pre-built `.onnx` inputs**. | | `--export-config` | | `PATH` | — | JSON ONNX export config overrides (opset version, constant folding, etc.) to merge into the Hugging Face export config. **Ignored for pre-built `.onnx` inputs**. | | `--dynamic-axes` | | `PATH` | — | JSON dynamic axes mapping for Hugging Face ONNX export, for example `{"input_ids": {"0": "batch", "1": "sequence"}}`. **Ignored for pre-built `.onnx` inputs**. | -| `--runtime` | | `winml-ort\|pytorch` | `winml-ort` | Evaluation runtime. `winml-ort` exports Hugging Face checkpoints to ONNX; `pytorch` evaluates the original checkpoint and supports `auto`, `cpu`, or CUDA-backed `gpu` devices. | +| `--runtime` | | `winml-ort\|winml-runtime\|pytorch` | `winml-ort` | Evaluation runtime. `winml-ort` exports Hugging Face checkpoints to ONNX; `winml-runtime` loads pre-built CGC MLIR; `pytorch` evaluates the original checkpoint and supports `auto`, `cpu`, or CUDA-backed `gpu` devices. | | `--dataset` | | `TEXT` | task default | HuggingFace dataset path (e.g., `imagenet-1k`, `nyu-mll/glue`). If omitted, a default dataset is selected based on the task. | | `--dataset-name` | | `TEXT` | — | Dataset configuration name for multi-config datasets. | | `--dataset-revision` | | `TEXT` | — | Git revision (branch, tag, or commit) of the dataset to load. Use `refs/convert/parquet` for HF datasets that are only served via the parquet mirror. | @@ -41,16 +41,16 @@ $ winml eval [options] | `--label-mapping` | | `PATH` | — | Path to a JSON file mapping dataset label names to the integer class IDs the model emits: `{"label_name": id}`. | | `--output` | `-o` | `PATH` | — | Output JSON file path for the evaluation results. | | `--schema` | | flag | `false` | Print the expected dataset schema for the given `--task` and exit. Does not run evaluation. | -| `--mode` | | `onnx\|compare` | `onnx` | Evaluation mode. `onnx` evaluates the ONNX candidate on a dataset. `compare` runs the ONNX candidate and a reference on identical random inputs and reports per-tensor similarity metrics — no dataset required. The reference is the HuggingFace model from `--model-id` by default, or a second ONNX file when `--reference` is given. | +| `--mode` | | `onnx\|compare` | `onnx` | Evaluation mode. `onnx` evaluates the candidate model on a dataset. `compare` runs the candidate and a reference on identical random inputs and reports per-tensor similarity metrics — no dataset required. The reference is the HuggingFace model from `--model-id` by default, or an ONNX file when `--reference` is given. | | `--input-data` | | `PATH` | — | Path to a `.npz` file of real input tensors to compare with instead of randomly generated ones (used with `--mode compare`). Keys must match the candidate model's input names. The **leading axis of each array is the sample axis**, so an archive shaped `(N, ...)` yields `N` samples (mean/std/min/max are computed across them); all inputs must share the same `N`. Each run is shaped to the candidate's batch size — a dynamic batch runs one row per sample, a static batch `B` chunks the axis into `N // B` batches (trailing rows are dropped with a warning). Note this differs from `winml perf --input-data`, which runs the **whole archive as a single batch**. | -| `--reference` | | `TEXT` | — | Reference `.onnx` file to compare the candidate against (used with `--mode compare`). Compares two ONNX models on identical random inputs; `--model-id` and `--task` are not required in this mode. | +| `--reference` | | `TEXT` | — | Reference `.onnx` file to compare against (used with `--mode compare`). `--model-id` and `--task` are not required in this mode. | | `--reference-device` | | `cpu\|gpu\|npu\|auto` | `cpu` | Device used for the reference ONNX model. Only valid with `--reference`. | | `--reference-device-luid` | | `TEXT` | — | Select a physical adapter for the reference ONNX model using its LUID from `winml sys`. Only valid with `--reference`. | | `--reference-ep` | | `TEXT` | — | Explicit execution provider used for the reference ONNX model, for example `dml`. Only valid with `--reference`. | ## How it works -`winml eval` loads the model and runs the evaluation pipeline via the internal `evaluate` function, then pulls the requested number of samples from a HuggingFace dataset. By default, Hugging Face model IDs and local checkpoints use the `winml-ort` runtime: they are exported to ONNX and evaluated through WinML. With `--runtime pytorch`, the task-resolved PyTorch model and stored dtype are preserved and the same dataset preprocessing, evaluator, and metrics run directly against that model. PyTorch `auto` selects CUDA when available and otherwise CPU; `gpu` requires CUDA. The JSON report identifies the effective runtime as `winml-ort` or `pytorch`. +`winml eval` loads the model and runs the evaluation pipeline via the internal `evaluate` function, then pulls the requested number of samples from a HuggingFace dataset. By default, Hugging Face model IDs and local checkpoints use the `winml-ort` runtime: they are exported to ONNX and evaluated through WinML. Pre-built CGC MLIR artifacts use `winml-runtime`. With `--runtime pytorch`, the task-resolved PyTorch model and stored dtype are preserved and the same dataset preprocessing, evaluator, and metrics run directly against that model. PyTorch `auto` selects CUDA when available and otherwise CPU; `gpu` requires CUDA. The JSON report identifies the effective runtime as `winml-ort`, `winml-runtime`, or `pytorch`. Python callers can pass an existing model directly with `evaluate(config, pytorch_model=model)`. An explicit `config.model_id` selects the tokenizer or processor; otherwise evaluation infers it from `model.config._name_or_path` and reports an error if neither source is available. diff --git a/docs/commands/optimize.md b/docs/commands/optimize.md index 5bd446a64..5c7134451 100644 --- a/docs/commands/optimize.md +++ b/docs/commands/optimize.md @@ -19,6 +19,7 @@ $ winml optimize [options] | `--model` | `-m` | `PATH` | *(required unless listing)* | Input ONNX model file. Not required when `--list-capabilities` or `--list-rewrites` is used. | | `--output` | `-o` | `PATH` | `{input}_opt.onnx` | Output path for the optimized model. Defaults to the input filename with `_opt` inserted before the extension. | | `--config` | `-c` | `PATH` | *(none)* | YAML or JSON configuration file. Fields in the file override capability defaults; CLI flags override the file. | +| `--enable-ort-graph-optimization` / `--disable-ort-graph-optimization` | | flag | enabled | Run or skip ORTGraphPipe. Disabling it does not implicitly enable compatibility rewrites or disable other pipes. Config key: `ort-graph-optimization` (boolean). | | `--verbose` | `-v` | flag | off | Enable verbose output. | | `--list-capabilities` | `-l` | flag | off | Print all registered optimization capabilities grouped by category and exit. Add `--verbose` for descriptions and ORT names. | | `--list-rewrites` | | flag | off | Print all available pattern-rewrite families with their source-to-target mappings and exit. | @@ -119,6 +120,75 @@ Enable any of the above with its --enable-* flag (dependencies are auto-enabled) Add `-v` to list every affected node and constant instead of a sample. +## CGC compatibility rewrites + +For an existing ONNX model, explicitly enable the required compatibility rules before +exporting CGIR MLIR. These rules are all **default off** and are independent of the +ordinary optimization/fusion defaults. + +CGC-specific patterns are isolated in `src\winml\modelkit\pattern\cgc`, with +private matching helpers in `cgc\utils.py` and public pattern exports in +`cgc\__init__.py`. The CGIR compatibility pipe owns rule selection and execution; +the shared pattern package does not re-export these backend-specific patterns. + +| Capability | Transformation and scope | +|------------|--------------------------| +| `normalize-int32-dq` | Normalize initializer-backed INT32 `DequantizeLinear` in the standard domain (opset >= 10) and `com.microsoft` (opset 1): omit immutable all-zero scalar/singleton zero points and clone singleton scales as scalars. Preserve shared initializers and domains; skip overridable parameters, per-axis vectors, unsupported attributes and nonlocal inputs. Handles nested graphs, not local functions. Enabled by CGC build configuration, disabled by default elsewhere. | +| `deduplicate-opset-imports` | Remove repeated model-level opset declarations with identical domain and version, retaining the first declaration and domain order. Reject conflicting versions for the same domain. Run before operator rewrites and opset upgrades; do not alter graph content, local functions, or the retained versions. | +| `eliminate-identity` | Remove safe internal tensor Identity aliases. Additionally replace top-level standard-domain FP32 graph-output Identities with same-shape Reshape when input/output types match exactly, all dimensions are positive static integers, opset >= 5, and the model has no subgraphs. Preserve output names/order, annotated aliases, unknown or conflicting types, scalar/dynamic/zero-size outputs and protected captures. Does not rewrite local functions. Workaround for [microsoft/ix#1198](https://github.com/microsoft/ix/issues/1198). | +| `gridsample-to-gather` | Decompose 2D `GridSample` with linear interpolation and zero padding into four `GatherND` reads, bounds masks and weighted sums. Supports both `align_corners` settings, FP16/FP32 IO and dynamic batch, with known positive channel, input spatial and grid spatial dimensions. Rank-3 indices contain explicit batch and spatial coordinates; `batch_dims=0` avoids the ORT symbolic shape inference defect tracked in [onnxruntime#24206](https://github.com/microsoft/onnxruntime/pull/24206). Batch coordinates are generated dynamically and shared across the four reads; sampled values are reshaped back to the grid layout. FP16 interpolation is computed in FP32 and cast back. Requires opset >= 16 (`bilinear` before opset 20); other modes are unchanged. Enabled by CGC builds, disabled in ordinary optimization. Floating-point rounding may differ from native sampling. | +| `omit-empty-resize-inputs` | Replace statically empty Resize ROI/scales with omitted inputs. Do not rely on graph-input defaults or rewrite crop-and-resize semantics. Requires opset 13; upgrade older matching models using ONNX version conversion. | +| `cgc-constant-folding` | Fill FoundryToolbox constant-folding gaps without an ORT Session. Fold standard `Pad.pads` constant integer chains; in graphs containing `Shape`, also fold statically known selected dimensions and bounded constant integer/boolean expressions (`Gather`, `Concat`, `Reshape`, `Slice`, `Transpose`, `Squeeze`, `Unsqueeze`, integer `Cast`, `ConstantOfShape`, arithmetic, `Equal`, `Where`). Iterate with shape inference, up to 32 rounds. Requires opset >= 11. Only the main graph is rewritten; preserve tensor names for shared uses and subgraph captures. Runtime floating-point computations and unresolved dimensions remain unchanged. This rule does not freeze inputs: specialize dimensions before optimization when needed; later Foundry `freeze-dims` does not retroactively affect this rule. Limits: 128 dependency values per traversal, 65,536 elements per operation and 1,048,576 cached elements per round. Enabled by CGC builds; disabled in ordinary optimization. `fold-constant-pad-pads` remains a compatibility alias. | +| `resize-tf-half-pixel-for-nn-to-asymmetric` | Change only the coordinate mode for nearest/floor Resize with static, non-overridable, positive integer scales. Dynamic/fractional scales and sizes-based inference are outside this rule. | +| `approximate-cubic-resize-with-linear` | **Lossy**, explicit cubic-to-linear approximation. Excludes antialiasing, outside exclusion, and crop-and-resize semantics. Prints a warning when applied. | +| `gathernd-to-reshape` | Replace GatherND only when data/indices/output ranks are not all equal and static, non-overridable int64 indices visit every input slice exactly once in storage order. Require positive static data dimensions; support batch dimensions, multi-coordinate indices, and equivalent negative indices. Dynamic data shapes or indices, overridable defaults, empty tensors, partial selection, repetition, and reordering are outside this rule. | +| `prelu-to-relu` | Decompose `PRelu(x, slope)` into `Relu(x) - slope * Relu(-x)` for FP16/FP32/FP64 and opset 7+. Require a direct, non-overridable initializer or Constant slope containing only finite values; preserve slope broadcasting and sharing. Dynamic data shapes are supported. Dynamic/non-finite slopes, integer types, and BF16 are outside this rule. | +| `dft-to-matmul` | Decompose DFT into real-valued sine/cosine matrix multiplications for FP16/FP32/FP64. Support opset 17+ axis attributes and opset 20+ scalar axis inputs; require known rank, static positive signal length, real/complex component count, and static optional axis/length parameters. Support forward/inverse full transforms, real forward onesided transforms, truncation, zero-padding, and dynamic batch dimensions. Dynamic transform parameters, BF16, and inverse onesided transforms are outside this rule. | + +Configuration keys can use underscores in place of hyphens. + +Add `--enable-deduplicate-opset-imports` (configuration: `"deduplicate_opset_imports": true`) +when repeated identical opset declarations leave a stale version during conversion. +This is an optimize rule, not an exporter fix: the saved optimized ONNX can be used +either for explicit CGIR export or directly by WinML Runtime's CGC backend. +It neither upgrades an opset nor resolves conflicting versions. + +```text +winml optimize -m model.onnx -o model.opt.onnx --disable-ort-graph-optimization --enable-omit-empty-resize-inputs --enable-resize-tf-half-pixel-for-nn-to-asymmetric +winml export -m model.opt.onnx -o model.mlir --target cgir +winml perf -m model.mlir --runtime winml-runtime --backend cgc --device gpu +``` + +Add `--enable-approximate-cubic-resize-with-linear` to optimize only when an accuracy-changing +approximation is acceptable. This option is not included in the example's preserving rewrites. + +Add `--enable-gathernd-to-reshape` (configuration: `"gathernd_to_reshape": true`) for +equivalent GatherND elimination. Rank mismatch alone is not sufficient: if identity +indexing cannot be proved, GatherND is retained. This bypasses CGC's DirectML descriptor +rank-alignment limitation; it does not implement general GatherND rank alignment. + +Add `--enable-prelu-to-relu` (configuration: `"prelu_to_relu": true`) to bypass IX's +missing ONNX PRelu legalization. CGC already has a parameterized-ReLU implementation; +the missing mapping is in IX. This is an algebraic decomposition, not an interpolation +approximation, but it does not guarantee preservation of the sign bit of zero. +Non-finite slopes are excluded because multiplying them by zero on the inactive +branch could introduce NaNs. + +Add `--enable-dft-to-matmul` (configuration: `"dft_to_matmul": true`) for DFT legalization. +The rewrite implements the DFT equation, including inverse normalization, without +requiring power-of-two lengths. It uses dense Fourier bases rather than an FFT: +constant storage and per-vector arithmetic are quadratic in the transform length. +It is intended as a compatibility workaround, not a performance replacement for a +native FFT. Different accumulation order and rounded coefficients mean floating-point +results are not bit-identical to the original DFT implementation. + +The CGIR compatibility pipe applies each enabled rule once, matching the result of the +preceding rule. This lets different rewrites affect the same Resize without changing +the shared rewrite pipe's conflict policy. Rewrites retain the original model IR version; +an opset upgrade may affect the whole model. These capabilities do not provide +general GatherND replacement, dynamic-length DFT decomposition, or general control-flow folding. +An upstream conversion/runtime limitation can still prevent export or execution. + ## Common pitfalls - **`--model` is required for actual optimization** — it can be omitted only when using `--list-capabilities` or `--list-rewrites`. Missing `--model` in any other case raises a usage error. diff --git a/docs/commands/perf.md b/docs/commands/perf.md index 8241f6c95..37c6efcea 100644 --- a/docs/commands/perf.md +++ b/docs/commands/perf.md @@ -1,6 +1,6 @@ # winml perf -> Benchmark an ONNX model's latency and throughput on a target device. +> Benchmark a model's latency and throughput on a target device. ## When to use this @@ -16,16 +16,16 @@ $ winml perf [options] | Flag | Short | Type | Default | Description | |---|---|---|---|---| -| `--model` | `-m` | `TEXT` | — | HuggingFace model ID or path to a local `.onnx` file. Required. With `--runtime ort-genai`, also accepts a prebuilt genai **bundle directory**, or a HuggingFace model ID that is auto-built into a bundle on demand. | -| `--runtime` | | `winml-ort\|ort-genai` | `winml-ort` | Inference runtime. `winml-ort` benchmarks single-shot ONNX inference; `ort-genai` benchmarks an onnxruntime-genai bundle (LLM generation: time-to-first-token + decode tokens/sec). With `ort-genai`, a model ID that is not a bundle directory is auto-built into one before benchmarking. An explicit `--ep` or `--device` selects both the transformer build and runtime target; without an override, the auto-build defaults to QNN/NPU. Bundles are cached under `~/.cache/winml/`, separately for each explicit EP/device target. GenAI cache controls are tracked in issue #1275. | +| `--model` | `-m` | `TEXT` | — | HuggingFace model ID or path to a local `.onnx` file. With `--runtime winml-runtime`, also accepts a prebuilt CGC `.mlir` file. Required. With `--runtime ort-genai`, also accepts a prebuilt genai **bundle directory**, or a HuggingFace model ID that is auto-built into a bundle on demand. | +| `--runtime` | | `auto\|winml-ort\|ort-genai\|winml-runtime` | `auto` | Inference runtime. `auto` selects `ort-genai` for local folders containing `genai_config.json`, `winml-runtime` for `.mlir` files, otherwise `winml-ort`; `winml-ort` benchmarks single-shot ONNX inference; `ort-genai` benchmarks an onnxruntime-genai bundle (LLM generation: time-to-first-token + decode tokens/sec); `winml-runtime` runs ONNX or prebuilt CGC MLIR through Windows ML Runtime. With `ort-genai`, a model ID that is not a bundle directory is auto-built into one before benchmarking. An explicit `--ep` or `--device` selects both the transformer build and runtime target; without an override, the auto-build defaults to QNN/NPU. Bundles are cached under `~/.cache/winml/`, separately for each explicit EP/device target. GenAI cache controls are tracked in issue #1275. | | `--task` | | `TEXT` | auto-detected | Explicit task override (e.g., `image-classification`). Inferred from the model if omitted. | | `--iterations` | | `INTEGER` | `100` (`10` with `--op-tracing`) | Number of timed inference iterations used to compute statistics. Explicit values override the op-tracing default. | | `--warmup` | | `INTEGER` | `10` | Number of warm-up iterations run before timing begins; excluded from statistics. | | `--device` | `-d` | `auto\|cpu\|gpu\|npu` | `auto` | Device to run the benchmark on. `auto` selects the highest-priority available device. | | `--device-luid` | | `TEXT` | — | Pin a physical adapter within the resolved EP/device pair using its LUID from `winml sys` (`0xHHHHHHHH_0xLLLLLLLL`, case-insensitive). Requires the EP to expose that adapter's LUID. Not supported with `--runtime ort-genai`. | | `--precision` | | `TEXT` | `auto` | Precision mode applied during model build: `auto`, `fp32`, `fp16`, `int8`, `int16`, or compound forms such as `w8a16`. | -| `--ep` | | `TEXT` | — | Force a specific execution provider (e.g., `qnn`, `dml`, `vitisai`, `openvino`, `cpu`). Overrides the device-to-provider mapping. | -| `--ep-options` | | `KEY=VALUE` (multiple) | — | Runtime EP provider option forwarded to the inference session (e.g., `--ep-options htp_performance_mode=burst`). Repeatable. Applies to both HuggingFace model IDs and ONNX file inputs. When detail op-tracing automatically compiles a raw ONNX model, these options are also applied to that compilation. | +| `--ep` | | `TEXT` | — | Force a specific execution provider (e.g., `qnn`, `dml`, `vitisai`, `openvino`, `cpu`). Overrides the device-to-provider mapping. With ONNX input and `--runtime winml-runtime`, the provider and `--device` class are passed to the Runtime execution target. Ignored for MLIR input. | +| `--ep-options` | | `KEY=VALUE` (multiple) | — | Runtime EP provider option forwarded to the inference session (e.g., `--ep-options htp_performance_mode=burst`). Repeatable. Applies to both HuggingFace model IDs and ONNX file inputs. When detail op-tracing automatically compiles a raw ONNX model, these options are also applied to that compilation. Ignored with `--runtime winml-runtime`. | | `--output` | `-o` | `PATH` | `~/.cache/winml/perf//.json` | Output JSON file path for the benchmark report. | | `--batch-size` | | `INTEGER` | `1` | Batch size used when generating synthetic input tensors. Ignored when `--input-data` is set. | | `--input-data` | | `PATH` | — | Path to a `.npz` file of real input tensors to benchmark with instead of randomly generated inputs. The archive's keys must match the model's inputs exactly; dtypes are cast to the model's expected dtype (with a warning) to mirror normal inference. Not supported with `--module`, `--runtime ort-genai`, or composite (dual-encoder) models. | @@ -167,6 +167,15 @@ Benchmark a pre-exported ONNX file on CPU with more iterations: $ winml perf -m model.onnx --device cpu --iterations 500 ``` +Benchmark a prebuilt CGC MLIR model with Windows ML Runtime: + +```bash +$ winml perf -m model.mlir --runtime winml-runtime --device gpu +``` + +For MLIR input, the resolved physical device is passed to Runtime as a DXCore +adapter target; the resolved EP is used only to identify that device. + Benchmark a text model with an explicit task, targeting the NPU: ```bash diff --git a/docs/getting-started/cgc-onboarding.md b/docs/getting-started/cgc-onboarding.md new file mode 100644 index 000000000..c154289f7 --- /dev/null +++ b/docs/getting-started/cgc-onboarding.md @@ -0,0 +1,168 @@ +# Preparing your model for the Windows ML Runtime APIs and DX CGC + +CGC onboarding has two journeys: + +1. **Benchmark the model:** Explore model compatibility and performance through the runtime + path that your application will use. +2. **Build the model:** Build a reusable `.mlir` model from PyTorch or ONNX. + +## Environment setup + +Create a virtual environment and install WinML CLI: + +```powershell +uv venv --python 3.11 +uv pip install winml-cli +``` + +Install the experimental wheels: + +```powershell +uv pip install "windowsml==" "onnxruntime-windowsml==" +``` +If you want to restore to stable wheel, you can run `uv sync` + + +Activate the virtual environment: + +```powershell +.venv\Scripts\Activate.ps1 +``` + +## Before you start + +Inspect the available devices and execution providers: + +```powershell +winml sys --list-device --list-ep +``` + +Verify that the WinMLCG EP is installed and that the target GPU and its +CGC-capable driver are listed. If they are missing, install or update the +Windows ML components and graphics driver before continuing. + +## Benchmark the model + +Use `winml perf` to explore model compatibility and performance before changing +your application. It exercises the same paths available to the application and +lets you compare them without first writing integration code. Use +[`winml eval`](../commands/eval.md) similarly to evaluate accuracy; the runtime +and EP options select the same CGC paths shown below. + +### MLIR with Runtime API + +Use this path when you already have an offline-converted `.mlir` model: + +```powershell +winml perf -m .\model.mlir --runtime winml-runtime -o .\model-mlir-perf.json +``` + +> **Note:** Currently, the graph and weights must be embedded in a single +> `.mlir` file. External weights will be supported in a future release. + +The Runtime API loads the model directly and runs it through DXCGC. A +`.mlir` model can only use the Windows ML Runtime API; it cannot be loaded by +ONNX Runtime and does not use an EP. + +### ONNX with WinMLCG EP + +Use this path when the application uses ONNX Runtime and should keep ONNX as +its model contract: + +```powershell +winml perf -m .\model.onnx --runtime winml-ort --ep winmlcg -o .\model-winmlcg-perf.json +``` + +The WinMLCG Execution Provider converts the model for CGC when the ONNX Runtime +session is created. No reusable CGIR artifact is produced. After this test +succeeds, keep the ONNX model and configure the WinMLCG EP in the application. + +### ONNX with Runtime API + +Use this path when the application uses the Windows ML Runtime API but should +keep ONNX as its model contract: + +```powershell +winml perf -m .\model.onnx --runtime winml-runtime -o .\model-runtime-perf.json +``` + +The CLI calls the Runtime compiler API to convert ONNX to temporary CGIR, +then loads that artifact and runs it through DXCGC. No EP is selected and no +reusable CGIR artifact is retained. Applications using this path must likewise +invoke the compiler API before pipeline execution. + +!!! warning + With `windowsml==2.7.25.dev0`, the tested CLI's online conversion path + fails with `Runtime.load_model() got an unexpected keyword argument 'io_counts'`. + The offline build and MLIR Runtime path below works with this combination. + +| Model input | Runtime path | EP | Conversion | +| --- | --- | --- | --- | +| `.onnx` | Windows ML ONNX Runtime | WinMLCG | When the session is created | +| `.onnx` | Windows ML Runtime API + DXCGC | None | When the session is created | +| `.mlir` | Windows ML Runtime API + DXCGC | None | Already converted offline | + +## Build the model + +Build a reusable `.mlir` model directly. No configuration file is required. +For PyTorch models, build first exports to ONNX. It then applies CGC +compatibility rewrites for supported operator patterns, quantizes the model, +and converts it to MLIR. + +> **Note:** By default, the CGC build converts FP32 to FP16. This may improve +> performance, but can reduce accuracy for some models. Add `--no-quant` to +> skip quantization if your model is already quantized or accuracy is a concern. + +### PyTorch → ONNX → CGIR + +Build from a Hugging Face model: + +```powershell +winml build -m microsoft/resnet-50 --backend cgc -o .\model-cgir +``` + +### ONNX → CGIR + +Build from an existing ONNX model: + +```powershell +winml build -m .\model.onnx --backend cgc -o .\model-cgir +``` + +Both paths write `model-cgir\model.onnx` and `model-cgir\model.mlir`. +Benchmark the `.mlir` model using the +[Windows ML Runtime API](#mlir-with-runtime-api). + +### Convert without optimization or quantization + +Use `winml export` to convert an ONNX model directly to MLIR, skipping the +optimization and quantization stages: + +```powershell +winml export -m .\model.onnx --target cgir -o .\model.mlir +``` + +The model must already be compatible with the CGC converter. + +## Tutorials and samples + +- [ResNet-50 - PyTorch to CGIR (experimental)](../samples/resnet50-cgir.md) walks through + the complete Hugging Face PyTorch → ONNX → CGIR build, performance, and + evaluation workflow. +- [YOLO11 - ONNX to CGIR (experimental)](../samples/yolo11-cgir.md) provides an end-to-end + conversion journey: export a checkpoint to ONNX, build CGIR, collect GPU + performance results, compare raw output tensors, and try the WinMLCG EP alternative. + +## Command references + +- [`winml export`](../commands/export.md) documents `--target cgir` and CGIR + export options. +- [`winml optimize`](../commands/optimize.md) documents ONNX optimization and + CGC compatibility rewrites. +- [`winml perf`](../commands/perf.md) documents `.mlir`, `winml-runtime`, and + WinMLCG performance measurement. +- [`winml eval`](../commands/eval.md) documents evaluation of supported CGIR + models with compatible preprocessing metadata. +- [`winml config`](../commands/config.md) and + [`winml build`](../commands/build.md) document configuration-driven CGIR + conversion. \ No newline at end of file diff --git a/docs/samples/resnet50-cgir.md b/docs/samples/resnet50-cgir.md new file mode 100644 index 000000000..dc44876bf --- /dev/null +++ b/docs/samples/resnet50-cgir.md @@ -0,0 +1,140 @@ +# ResNet-50 - PyTorch to CGIR (experimental) + +This tutorial starts with the +[`microsoft/resnet-50`](https://huggingface.co/microsoft/resnet-50) PyTorch +model on Hugging Face, exports it to ONNX, converts it to a reusable CGIR +`.mlir` artifact, and measures it through the Windows ML Runtime API. + +You will use a generated build configuration so the complete PyTorch → ONNX → +CGIR workflow can be reviewed, repeated, and checked into source control. + +## Prerequisites + +- A Windows device with a GPU. +- An activated Python 3.11 environment with winml-cli and the experimental + wheels from [CGC environment setup](../getting-started/cgc-onboarding.md#environment-setup). +- A network connection to download the model and evaluation dataset. + +## Step 1: Check the CGC environment + +Inspect the devices and execution providers available on your machine: + +```powershell +winml sys --list-device --list-ep +``` + +Verify that the WinMLCG EP and target GPU are listed. Discovery alone does +not prove that the driver supports CGC execution; the benchmark below checks +that path on your device. + +## Step 2: Generate a CGC build config + +Generate a reusable configuration for the Hugging Face model: + +```powershell +winml config -m microsoft/resnet-50 --backend cgc -o .\resnet50-cgir\config.json +``` + +`--backend cgc` configures the pipeline to export the PyTorch model to ONNX, +prepare it for CGC, and convert the final ONNX model to CGIR. The generated +configuration uses FP16 conversion and does not run a separate EP compilation +stage. + +Open `resnet50-cgir\config.json` if you want to review the detected +image-classification task, model loader, ONNX export settings, and +`convert.target: "cgir"` before building. + +## Step 3: Build ONNX and CGIR artifacts + +Run the configured build: + +```powershell +winml build -m microsoft/resnet-50 -c .\resnet50-cgir\config.json -o .\resnet50-cgir\build +``` + +The first run downloads the model from Hugging Face. The build exports the +PyTorch model to ONNX, applies the configured transformations, and writes the +CGIR model. The primary artifacts are: + +```text +resnet50-cgir/ +├── config.json +└── build/ + ├── model.onnx + └── model.mlir +``` + +The build directory can contain additional reports and intermediate files. +Keep any weight sidecars beside `model.mlir` when moving the CGIR model. + +### How the pipeline works + +`winml config` does not transform the model. It detects the Hugging Face task, +model class, and ONNX export settings, then writes the pipeline configuration +consumed by `winml build`. For `--backend cgc`, that configuration enables FP16 +preparation, disables EP compilation, and adds `convert.target: "cgir"`. + +`winml build` combines these primitive stages: + +1. `winml export` converts the PyTorch model to ONNX. +2. `winml optimize` rewrites ONNX operators to better align with the current + CGC IR. +3. `winml quantize` converts the model from FP32 to FP16. +4. The CGIR export stage converts the optimized ONNX model to `.mlir`. + +FP16 conversion can affect model accuracy. If accuracy drops, rebuild with +`--no-quant` and compare the results. + +!!! important + The `.mlir` artifact cannot be used directly by ONNX Runtime. Load it + through the Windows ML Runtime API for CGC inference. + +## Step 4: Measure CGIR performance + +Run the CGIR artifact through the Windows ML Runtime API and save the +performance results: + +```powershell +winml perf -m .\resnet50-cgir\build\model.mlir --runtime winml-runtime --output .\resnet50-cgir\perf.json +``` + +The report includes warm-up and timed iteration counts, latency percentiles, +throughput, and the resolved device. For `.mlir` input, `winml perf` +automatically selects `winml-runtime`; the explicit option documents the +intended runtime path. + +## Step 5: Evaluate model accuracy + +First measure the source PyTorch model on a fixed sample from mini-ImageNet: + +```powershell +winml eval -m microsoft/resnet-50 --runtime pytorch --device cpu --dataset timm/mini-imagenet --split test --samples 100 --no-shuffle -o .\resnet50-cgir\pytorch-eval.json +``` + +Evaluate the CGIR artifact on the same samples. `--model-id` provides the +Hugging Face image processor and label configuration associated with the local +`.mlir` model: + +```powershell +winml eval -m .\resnet50-cgir\build\model.mlir --model-id microsoft/resnet-50 --runtime winml-runtime --dataset timm/mini-imagenet --split test --samples 100 --no-shuffle -o .\resnet50-cgir\cgir-eval.json +``` + +Compare the accuracy in `pytorch-eval.json` and `cgir-eval.json`. Both commands +use the same ordered 100-sample slice, so a difference reflects the exported +runtime path rather than a different random sample. + +## What you built + +- `config.json`: the repeatable CGC build configuration. +- `build\model.onnx`: the optimized, FP16-prepared ONNX model used for CGIR conversion. +- `build\model.mlir`: the reusable CGIR model for Windows ML Runtime. +- `perf.json`: GPU latency and throughput results. +- `pytorch-eval.json` and `cgir-eval.json`: source and CGIR accuracy results. + +## See also + +- [CGC Onboarding (experimental)](../getting-started/cgc-onboarding.md) +- [winml config](../commands/config.md) +- [winml build](../commands/build.md) +- [winml perf](../commands/perf.md) +- [winml eval](../commands/eval.md) \ No newline at end of file diff --git a/docs/samples/yolo11-cgir.md b/docs/samples/yolo11-cgir.md new file mode 100644 index 000000000..623f28d76 --- /dev/null +++ b/docs/samples/yolo11-cgir.md @@ -0,0 +1,106 @@ +# YOLO11 - ONNX to CGIR (experimental) + +This sample starts with an Ultralytics YOLO11 object-detection checkpoint, +exports an ONNX model, converts that model to a reusable CGIR artifact, and +measures it on the GPU. It also shows the ONNX + WinMLCG Execution Provider +path for testing CGC without keeping a CGIR file. + +## Before you start + +Complete [CGC environment setup](../getting-started/cgc-onboarding.md#environment-setup) +and activate that virtual environment. Inspect the available EPs and devices: + +```powershell +winml sys --list-device --list-ep +``` + +## Step 1: Export YOLO11 to ONNX + +Install the exporter into the same environment: + +```powershell +uv pip install "ultralytics==8.4.146" --index-url https://packagefeedproxy.microsoft.io/pypi/simple +``` + +Place the Ultralytics `yolo11n.pt` checkpoint in the current directory, then +export a static, batch-one FP32 model without ONNX simplification: + +```powershell +yolo export model=yolo11n.pt format=onnx imgsz=640 batch=1 dynamic=False simplify=False opset=17 device=cpu +``` + +This produces `yolo11n.onnx` with input shape `[1, 3, 640, 640]` and raw +detection output shape `[1, 84, 8400]`. + +## Step 2: Benchmark ONNX + +Test ONNX through the WinMLCG EP: + +```powershell +winml perf -m .\yolo11n.onnx --runtime winml-ort --ep winmlcg -o .\model-winmlcg-perf.json +``` + +!!! warning + With `windowsml==2.7.25.dev0` and `onnxruntime-windowsml==1.30.0.202609102321`, + two full benchmark runs produced reports but exited with code 1 on the + validation machine. The cause is unresolved; a saved report alone does + not indicate a successful command. + +To test online CGIR conversion through the Windows ML Runtime API instead: + +```powershell +winml perf -m .\yolo11n.onnx --runtime winml-runtime -o .\model-runtime-perf.json +``` + +!!! warning + With `windowsml==2.7.25.dev0`, the tested CLI's online conversion path + fails with `Runtime.load_model() got an unexpected keyword argument 'io_counts'`. + Use the offline build and MLIR Runtime path below with this combination. + +## Step 3: Build CGIR + +Generate the CGC build configuration: + +```powershell +winml config -m .\yolo11n.onnx --backend cgc -o .\cgc-config.json +``` + +Build the reusable `.mlir` model: + +```powershell +winml build -m .\yolo11n.onnx -c .\cgc-config.json -o .\model-cgir +``` + +The default configuration applies CGC compatibility rewrites and FP16 +conversion before exporting MLIR. FP16 can change accuracy; add `--no-quant` +to the build command to skip quantization. + +## Step 4: Benchmark CGIR + +Run the converted model through the Windows ML Runtime API: + +```powershell +winml perf -m .\model-cgir\model.mlir --runtime winml-runtime -o .\model-cgir\model-perf.json +``` + +!!! important + The `.mlir` artifact cannot be used directly by ONNX Runtime. Load it + through the Windows ML Runtime API for CGC inference. + +## Step 5: Compare output tensors + +Compare the CGIR model against the original FP32 ONNX model on ten generated +inputs, using ONNX Runtime on CPU as the reference: + +```powershell +winml eval --mode compare -m .\model-cgir\model.mlir --runtime winml-runtime --reference .\yolo11n.onnx --reference-device cpu --samples 10 -o .\model-cgir\tensor-similarity.json +``` + +This checks raw output similarity, not object-detection accuracy. Inspect +absolute errors as well as cosine similarity. Detection mAP still requires +a labeled dataset with matching preprocessing and postprocessing and has not +been validated by this sample. + +## See also + +- [CGC Onboarding](../getting-started/cgc-onboarding.md) diff --git a/src/winml/modelkit/build/hf.py b/src/winml/modelkit/build/hf.py index 4094b8938..80cc20ee4 100644 --- a/src/winml/modelkit/build/hf.py +++ b/src/winml/modelkit/build/hf.py @@ -21,6 +21,7 @@ import datetime import gc +import json import logging import time from dataclasses import dataclass, field @@ -28,8 +29,9 @@ from typing import TYPE_CHECKING, Any from ..export import export_onnx +from ..onnx import copy_onnx_model from ..utils import MANIFEST_FILENAME, ManifestStage, WinMLManifest -from .common import run_build_stages +from .common import StagesResult, run_build_stages if TYPE_CHECKING: @@ -83,6 +85,7 @@ def build_hf_model( model_id: str | None = None, pytorch_model: nn.Module | None = None, rebuild: bool = False, + skip_build: bool = False, trust_remote_code: bool = False, random_init: bool = False, cache_key: str | None = None, @@ -111,6 +114,7 @@ def build_hf_model( pytorch_model: Pre-loaded PyTorch model. If provided, model_id is only used for labeling (not loading). rebuild: If True, overwrite existing artifacts and re-run pipeline. + skip_build: Export ONNX without optimization, analysis, quantization, or compilation. trust_remote_code: Whether to trust remote code when loading HF models. cache_key: Optional prefix for artifact filenames. ep: Target execution provider for the analyzer (e.g., ``"qnn"``). @@ -241,22 +245,31 @@ def _name(base: str) -> str: # Shared with build_onnx_model via ``common.run_build_stages``. # ========================================================================= skip_optimize: bool = kwargs.pop("skip_optimize", False) - stages = run_build_stages( - current_path=export_path, - optimized_path=optimized_path, - quantized_path=quantized_path, - compiled_path=compiled_path, - final_path=final_path, - config=config, - config_path=config_path, - ep=ep, - device=device, - hack_max_optim_iterations=hack_max_optim_iterations, - skip_optimize=skip_optimize or config.skip_optimize, - allow_unsupported_nodes=allow_unsupported_nodes, - analyze_result_path=output_dir / _name("analyze_result.json"), - onnx_kwargs=onnx_kwargs, - ) + if skip_build: + copy_onnx_model(export_path, final_path) + config_path.write_text(json.dumps(config.to_dict(), indent=2)) + stages = StagesResult( + current_path=final_path, + is_pre_quantized=False, + stages_skipped=["optimize", "quantize", "compile"], + ) + else: + stages = run_build_stages( + current_path=export_path, + optimized_path=optimized_path, + quantized_path=quantized_path, + compiled_path=compiled_path, + final_path=final_path, + config=config, + config_path=config_path, + ep=ep, + device=device, + hack_max_optim_iterations=hack_max_optim_iterations, + skip_optimize=skip_optimize or config.skip_optimize, + allow_unsupported_nodes=allow_unsupported_nodes, + analyze_result_path=output_dir / _name("analyze_result.json"), + onnx_kwargs=onnx_kwargs, + ) stages_completed.extend(stages.stages_completed) stages_skipped.extend(stages.stages_skipped) stage_timings.update(stages.stage_timings) diff --git a/src/winml/modelkit/commands/_pre_bench.py b/src/winml/modelkit/commands/_pre_bench.py index 8a1f5f3fe..15ce8da2f 100644 --- a/src/winml/modelkit/commands/_pre_bench.py +++ b/src/winml/modelkit/commands/_pre_bench.py @@ -47,6 +47,7 @@ def print_pre_bench_block( ep_source: str, ep_version: str | None, ep_dll_path: str, + runtime_api_backend: str | None = None, ) -> None: """Print the pre-benchmark identity block. @@ -84,6 +85,8 @@ def print_pre_bench_block( ``v`` chunk when absent). ep_dll_path: Full path to the plugin DLL. Empty string signals a built-in EP and renders as ``(bundled with ORT)``. + runtime_api_backend: Runtime API backend, when applicable. CGC does + not expose EP information. """ # --- Model panel: identity + surface --------------------------------- model_lines: list[Text] = [] @@ -112,16 +115,20 @@ def print_pre_bench_block( # --- Device panel: resolved device + EP + DLL ------------------------- hw_suffix = f" [dim]({hardware_name})[/dim]" if hardware_name else "" - ep_line = f"[cyan]{ep}[/cyan]@[cyan]{ep_source}[/cyan]" - if ep_version: - ep_line += f" [green]v{ep_version}[/green]" - dll_display = ep_dll_path if ep_dll_path else "(bundled with ORT)" - device_lines: list[Text] = [ _labeled_line("Device:", f"[cyan]{device}[/cyan]{hw_suffix}"), - _labeled_line("EP:", ep_line), - _labeled_line("EP DLL:", f"[dim]{dll_display}[/dim]"), ] + if runtime_api_backend != "cgc": + ep_line = f"[cyan]{ep}[/cyan]@[cyan]{ep_source}[/cyan]" + if ep_version: + ep_line += f" [green]v{ep_version}[/green]" + dll_display = ep_dll_path if ep_dll_path else "(bundled with ORT)" + device_lines.extend( + [ + _labeled_line("EP:", ep_line), + _labeled_line("EP DLL:", f"[dim]{dll_display}[/dim]"), + ] + ) console.print(Panel(Group(*device_lines), title="Device", expand=True)) diff --git a/src/winml/modelkit/commands/build.py b/src/winml/modelkit/commands/build.py index a56d9f52e..ceac3fb8b 100644 --- a/src/winml/modelkit/commands/build.py +++ b/src/winml/modelkit/commands/build.py @@ -37,6 +37,7 @@ print_setup, print_stages_header, ) +from ..utils.constants import RUNTIME_BACKENDS, RuntimeBackend from ..utils.logging import configure_logging from ..utils.model_input import ModelInputKind, classify_model_input from ._ep_arg import EpAtSourceParamType @@ -836,6 +837,12 @@ def _maybe_build_genai_bundle( @cli_utils.precision_option( optional_message="With -c, applied only when --device or --precision is passed.", ) +@click.option( + "--backend", + type=click.Choice(list(RUNTIME_BACKENDS)), + default=None, + help="Backend used when auto-generating config, as in winml config --backend.", +) @click.option( "--export-type", type=click.Choice(["generic", "optimized"], case_sensitive=False), @@ -904,6 +911,7 @@ def build( submodel: str | None, verbose: int, quiet: bool, + backend: RuntimeBackend | None = None, ) -> None: r"""Build a WinML-optimized ONNX model from a HuggingFace model or .onnx file. @@ -918,6 +926,9 @@ def build( # Auto-generate config (no -c needed) winml build -m microsoft/resnet-50 -o output/ + # One-step CGIR build, preserving the existing ONNX precision + winml build -m model.onnx -o output/ --backend cgc --no-quant + # Full pipeline with explicit config winml build -c config.json -m microsoft/resnet-50 -o output/ @@ -950,6 +961,9 @@ def build( # are needed, or pin --ep qnn --device npu to build it on any host) winml build -m Qwen/Qwen3-0.6B -o out/ --export-type optimized """ + if backend == "cgc" and ep is not None: + raise click.UsageError("--backend cgc cannot be combined with --ep.") + # Merge top-level -v/-q with subcommand-level flags so either position works. verbose, quiet = cli_utils.resolve_verbosity(ctx, verbose, quiet) configure_logging(verbosity=verbose, quiet=quiet) @@ -1087,6 +1101,7 @@ def build( device=runtime_device, precision=precision, ep=runtime_ep_value, + backend=backend, ) else: config_or_configs = generate_build_config( @@ -1095,6 +1110,7 @@ def build( device=runtime_device, precision=precision, ep=runtime_ep_value, + backend=backend, export_policy_target=(request_device, request_ep_value), shape_config=shape_overrides, override={"export": export_overrides} if export_overrides else None, @@ -1125,7 +1141,7 @@ def _patch_device(cfg: WinMLBuildConfig) -> None: from ..config import resolve_quant_compile_config resolved_quant, _ = resolve_quant_compile_config( - device=runtime_device, precision=precision, ep=runtime_ep_value + device=runtime_device, precision=precision, ep=runtime_ep_value, backend=backend ) if not quant or resolved_quant is None or is_pre_quantized_onnx_input: cfg.quant = None @@ -1419,6 +1435,7 @@ def _patch_device(cfg: WinMLBuildConfig) -> None: device=runtime_device, precision=precision, ep=runtime_ep_value, + backend=backend, export_policy_target=(request_device, request_ep_value), shape_config=shape_overrides, override={"export": export_overrides} if export_overrides else None, @@ -1464,6 +1481,7 @@ def _patch_device(cfg: WinMLBuildConfig) -> None: component_config.compile = None else: component_config.compile = copy.deepcopy(config.compile) + component_config.convert = copy.deepcopy(config.convert) try: component_config.validate() @@ -1627,9 +1645,15 @@ def _run_single_build( preloaded_hf_config=preloaded_hf_config, ) - elapsed = time.monotonic() - start_time final_name = f"{cache_key}_model.onnx" if cache_key else "model.onnx" final_path = resolved_dir / final_name + stage_timings = stage_timings or [] + final_path = _run_convert_stage( + config=config, + current_path=final_path, + stage_timings=stage_timings, + ) + elapsed = time.monotonic() - start_time if final_path.exists() and stage_timings: config_json = resolved_dir / ( f"{cache_key}_winml_build_config.json" if cache_key else "winml_build_config.json" @@ -2039,6 +2063,35 @@ def _run_compile_stage( return current_path +def _run_convert_stage( + *, + config: WinMLBuildConfig, + current_path: Path, + stage_timings: list[tuple[str, float | None]], +) -> Path: + """Convert the final ONNX artifact to MLIR when configured.""" + if config.convert is None: + return current_path + + from ..export.cgc import CGCExporter + from ..utils.console import StageLive + + options = cli_utils.parse_options( + tuple(f"{key}={value}" for key, value in config.convert.options.items()), + CGCExporter.options_type, + ) + output_path = current_path.with_suffix(".mlir") + with StageLive("convert", console) as sl: + sl.set_status("Converting ONNX to MLIR...") + started = time.monotonic() + CGCExporter(options).export_onnx(model=current_path, output_path=output_path) + elapsed = time.monotonic() - started + sl.set_done(elapsed) + sl.artifact(str(output_path), _safe_size(output_path)) + stage_timings.append(("Convert", elapsed)) + return output_path + + # ============================================================================= # PIPELINE FUNCTIONS # ============================================================================= @@ -2205,7 +2258,7 @@ def _build_onnx_pipeline( or None if build was reused. """ from ..build.common import ensure_pre_quantized_stamped - from ..onnx import copy_onnx_model + from ..onnx import copy_onnx_model, is_quantized_onnx max_iters: int = extra_kwargs.pop("hack_max_optim_iterations", 3) allow_unsupported_nodes: bool = extra_kwargs.pop("allow_unsupported_nodes", False) @@ -2246,11 +2299,14 @@ def _build_onnx_pipeline( if current_path.resolve() != onnx_path.resolve(): copy_onnx_model(onnx_path, current_path) - # Keep the CLI ONNX path aligned with the library build paths: if a user - # supplies a pre-quantized model via ``-c config.json`` we must stamp the - # config before any stage reads it, otherwise the optimize stage will still - # run on integer ops and the quantize stage may try to re-quantize. - ensure_pre_quantized_stamped(config, current_path) + if not config.is_cgc: + # Keep the CLI ONNX path aligned with the library build paths: if a user + # supplies a pre-quantized model via ``-c config.json`` we must stamp the + # config before any stage reads it, otherwise the optimize stage will still + # run on integer ops and the quantize stage may try to re-quantize. + ensure_pre_quantized_stamped(config, current_path) + elif config.quant is not None and is_quantized_onnx(current_path): + config.quant = None # ── Optimize stage (first stage for ONNX — show I/O here) ──── current_path, _ = _run_optimize_stage( diff --git a/src/winml/modelkit/commands/config.py b/src/winml/modelkit/commands/config.py index 60317e09f..46baa6b21 100644 --- a/src/winml/modelkit/commands/config.py +++ b/src/winml/modelkit/commands/config.py @@ -40,6 +40,7 @@ print_kv, print_success, ) +from ..utils.constants import RUNTIME_BACKENDS, RuntimeBackend from ..utils.logging import configure_logging from ..utils.model_input import ModelInputKind, classify_model_input from ._ep_arg import EpAtSourceParamType @@ -147,6 +148,13 @@ def _merge_export_overrides(cfg: Any, export_overrides: dict[str, Any]) -> Any: "takes a bare EP short-name.)", ) @cli_utils.precision_option() +@click.option( + "--backend", + type=click.Choice(list(RUNTIME_BACKENDS)), + default=None, + help="Build backend. cgc generates FP16 and CGIR conversion without compilation; " + "ort or omission preserves the existing configuration behavior.", +) @cli_utils.output_option("Output JSON file path (default: stdout)") @cli_utils.overwrite_option() @click.option( @@ -190,6 +198,7 @@ def config( quant: bool, no_compile: bool, trust_remote_code: bool, + backend: RuntimeBackend | None = None, ) -> None: r"""Generate WinMLBuildConfig for a HuggingFace model or .onnx file. @@ -239,6 +248,9 @@ def config( # Generate configs for submodules winml config -m microsoft/resnet-50 --module ResNetConvLayer """ + if backend == "cgc" and ep is not None: + raise click.UsageError("--backend cgc cannot be combined with --ep.") + verbose, quiet = cli_utils.resolve_verbosity(ctx, verbose, quiet) configure_logging(verbosity=verbose, quiet=quiet) @@ -364,6 +376,7 @@ def config( device=device, precision=precision, ep=ep_name, + backend=backend, override=onnx_override, ) @@ -414,6 +427,7 @@ def config( ep=ep_name, no_quant=not quant, no_compile=no_compile, + backend=backend, policy_overrides_config=policy_overrides_config, output=output, overwrite=overwrite, @@ -447,6 +461,7 @@ def config( precision=precision, trust_remote_code=trust_remote_code, ep=ep_name, + backend=backend, policy_overrides_config=policy_overrides_config, ) if isinstance(result, list): @@ -650,6 +665,7 @@ def _generate_pipeline_configs( output: Path | None, overwrite: bool, console: Any, + backend: RuntimeBackend | None = None, ) -> None: """Generate and save one config file per pipeline sub-component.""" from ..config import generate_hf_build_config @@ -672,6 +688,7 @@ def _generate_pipeline_configs( precision=precision, trust_remote_code=trust_remote_code, ep=ep, + backend=backend, policy_overrides_config=policy_overrides_config, ) _apply_stage_overrides(cfg, no_quant=no_quant, no_compile=no_compile) diff --git a/src/winml/modelkit/commands/eval.py b/src/winml/modelkit/commands/eval.py index 261f2c13b..16a866958 100644 --- a/src/winml/modelkit/commands/eval.py +++ b/src/winml/modelkit/commands/eval.py @@ -17,7 +17,12 @@ from rich.console import Console from ..utils import cli as cli_utils -from ..utils.constants import ALL_EP_NAMES, SUPPORTED_DEVICES +from ..utils.constants import ( + ALL_EP_NAMES, + RUNTIME_BACKENDS, + SUPPORTED_DEVICES, + resolve_runtime_api_backend, +) from ..utils.eval_utils import EVAL_MODES, TASK_SCHEMAS, EvalMode, TaskSchema from ..utils.logging import configure_logging @@ -35,12 +40,12 @@ required=False, multiple=True, help_text=( - "Model to evaluate. Accepts a HuggingFace model ID, an ONNX file path " + "Model to evaluate. Accepts a HuggingFace model ID, an ONNX or MLIR file path " "(requires --model-id), or split-encoder role=path pairs (see --schema)." ), ) @cli_utils.model_id_option( - help_text="HuggingFace model ID when .onnx model file is provided in --model.", + help_text="HuggingFace model ID when an ONNX or MLIR file is provided in --model.", ) @click.option( "--dataset", @@ -119,11 +124,18 @@ ) @click.option( "--runtime", - type=click.Choice(["winml-ort", "pytorch"]), + type=click.Choice(["winml-ort", "winml-runtime", "pytorch"]), default="winml-ort", show_default=True, help="Evaluation runtime. 'winml-ort' exports Hugging Face checkpoints to ONNX; " - "'pytorch' evaluates the original checkpoint.", + "'winml-runtime' loads pre-built MLIR; 'pytorch' evaluates the original checkpoint.", +) +@click.option( + "--backend", + type=click.Choice(list(RUNTIME_BACKENDS)), + default=None, + help="[winml-runtime] Execution backend for ONNX inputs (default: cgc). " + "MLIR inputs always use cgc.", ) @click.option( "--samples", @@ -194,8 +206,8 @@ show_default=True, help=( "Evaluation mode. " - "'onnx' (default): evaluate the ONNX candidate on the dataset. " - "'compare': compare ONNX vs HF reference output tensors on identical " + "'onnx' (default): evaluate the candidate model on the dataset. " + "'compare': compare candidate vs reference output tensors on identical " "random inputs and report tensor-similarity metrics per output tensor." ), ) @@ -218,7 +230,7 @@ default=None, help=( "Reference ONNX file to compare the candidate against (use with " - "--mode compare). Compares two ONNX models on identical random inputs; " + "--mode compare). Compares two models on identical random inputs; " "--model-id / --task are not required in this mode." ), ) @@ -269,6 +281,7 @@ def eval( export_config: Path | None, dynamic_axes: Path | None, runtime: EvalRuntime, + backend: str | None, ep: EPNameOrAlias | None, samples: int, split: str, @@ -339,9 +352,10 @@ def eval( # ── 1. Build config: defaults ← config file ← CLI ── cfg, config_fields = _build_eval_config(ctx, config_file, column, label_mapping_path) - if cfg.runtime not in ("winml-ort", "pytorch"): + if cfg.runtime not in ("winml-ort", "winml-runtime", "pytorch"): raise click.UsageError( - f"Invalid eval runtime {cfg.runtime!r}; expected 'winml-ort' or 'pytorch'." + f"Invalid eval runtime {cfg.runtime!r}; expected 'winml-ort', " + "'winml-runtime', or 'pytorch'." ) if cfg.runtime == "pytorch": _validate_pytorch_runtime_options(ctx, cfg, config_fields) @@ -368,6 +382,16 @@ def eval( # ── 2. Resolve in place ── _resolve_model(cfg, model, model_id, allow_missing_model_id=cfg.reference_path is not None) + is_mlir = ( + isinstance(cfg.model_path, str) + and Path(cfg.model_path).suffix.lower() == ".mlir" + ) + if is_mlir and cfg.runtime != "winml-runtime": + raise click.UsageError("MLIR inputs require --runtime winml-runtime.") + try: + cfg.backend = resolve_runtime_api_backend(cfg.runtime, cfg.model_path, cfg.backend) + except ValueError as error: + raise click.UsageError(str(error)) from error if cfg.runtime == "pytorch" and cfg.model_path is not None: raise click.UsageError( "--runtime pytorch requires a Hugging Face model ID or local Hugging Face " @@ -668,9 +692,9 @@ def _resolve_model( def _resolve_reference(cfg: WinMLEvaluationConfig) -> None: - """Validate and normalize ``cfg.reference_path`` for two-ONNX compare. + """Validate and normalize ``cfg.reference_path`` for ONNX compare. - Requires the candidate (``-m``) to be a single ONNX file (composite + Requires the candidate (``-m``) to be a single model file (composite ``role=path`` candidates and build-from-id are not supported with ``--reference`` yet). Resolves Hub-hosted ONNX refs to local paths. """ @@ -679,7 +703,7 @@ def _resolve_reference(cfg: WinMLEvaluationConfig) -> None: if not isinstance(cfg.model_path, str): raise click.UsageError( - "--reference requires the candidate (-m) to be a single ONNX file. " + "--reference requires the candidate (-m) to be a single model file. " "Composite (role=path) candidates and build-from-id are not " "supported with --reference." ) @@ -780,7 +804,8 @@ def _resolve_device(cfg: WinMLEvaluationConfig) -> None: console = Console(stderr=True) console.print("[bold]Detecting available devices...[/bold]") resolved_target = resolve_device( - EPDeviceTarget(ep=cfg.ep or "auto", device=cfg.device or "auto") + EPDeviceTarget(ep=cfg.ep or "auto", device=cfg.device or "auto"), + backend=resolve_runtime_api_backend(cfg.runtime, cfg.model_path, cfg.backend), ) cfg.device = resolved_target.device console.print(f"[dim]Using device:[/dim] {resolved_target.device}") @@ -953,9 +978,9 @@ def _resolve_model_path( ) value = plain[0] - if Path(value).suffix.lower() == ".onnx": - # Hub-hosted ONNX (e.g. ``onnx-community/sam3-tracker-ONNX/onnx/...``) - # is downloaded once and treated as a local .onnx path thereafter. + model_suffix = Path(value).suffix.lower() + if model_suffix in (".onnx", ".mlir"): + # Hub-hosted artifacts are resolved once; local paths pass through. try: value = cli_utils.normalize_model_arg(value) or value except Exception as e: @@ -964,14 +989,14 @@ def _resolve_model_path( ) from e if not Path(value).exists(): raise click.BadParameter( - f"ONNX file not found: {value}", + f"{model_suffix.removeprefix('.').upper()} file not found: {value}", param_hint="-m/--model", ) if model_id is None: if allow_missing_model_id: return value, None raise click.UsageError( - "When using an ONNX file, --model-id is required " + f"When using a {model_suffix} file, --model-id is required " "for preprocessor and config resolution." ) return value, model_id @@ -988,7 +1013,7 @@ def _resolve_model_path( if model_id is not None and model_id != value: raise click.UsageError( "Cannot pass both `-m ` and `--model-id`. " - "Use `--model-id` only together with an ONNX file path in `-m`." + "Use `--model-id` only together with an ONNX or MLIR file path in `-m`." ) return None, model_id or value @@ -1014,6 +1039,7 @@ def display_eval_report(result: EvalResult, console: Console) -> None: cfg = result.config ds = cfg.dataset metrics = result.metrics + backend = resolve_runtime_api_backend(cfg.runtime, cfg.model_path, cfg.backend) # For --input-data compare the effective sample count comes from the # archive (via EvalResult.num_samples), not the unused config default. samples = result.num_samples if result.num_samples is not None else ds.samples @@ -1039,23 +1065,38 @@ def display_eval_report(result: EvalResult, console: Console) -> None: # Info section console.print() console.print(f"[dim]Task:[/dim] {cfg.task}") - console.print(f"[dim]Runtime:[/dim] {cfg.runtime}") - console.print(f"[dim]Device:[/dim] {cfg.device}") if cfg.input_data: console.print(f"[dim]Input data:[/dim] {cfg.input_data}") elif ds.path: console.print(f"[dim]Dataset:[/dim] {ds.path}") console.print(f"[dim]Samples:[/dim] {samples}") - if isinstance(cfg.model_path, dict): + if cfg.mode == "compare": + console.print(f"[dim]Candidate:[/dim] {cfg.model_path or cfg.model_id}") + console.print(f"[dim]Candidate runtime:[/dim] {cfg.runtime}") + console.print(f"[dim]Candidate device:[/dim] {cfg.device}") + if backend != "cgc": + console.print(f"[dim]Candidate EP:[/dim] {cfg.ep or 'auto'}") + console.print(f"[dim]Reference:[/dim] {cfg.reference_path or cfg.model_id}") + console.print( + f"[dim]Reference runtime:[/dim] " + f"{'winml-ort' if cfg.reference_path else 'pytorch'}" + ) + console.print( + f"[dim]Reference device:[/dim] " + f"{cfg.reference_device if cfg.reference_path else 'cpu'}" + ) + console.print( + f"[dim]Reference EP:[/dim] " + f"{(cfg.reference_ep or 'auto') if cfg.reference_path else 'n/a'}" + ) + else: + console.print(f"[dim]Runtime:[/dim] {cfg.runtime}") + console.print(f"[dim]Device:[/dim] {cfg.device}") + if cfg.mode != "compare" and isinstance(cfg.model_path, dict): for role, path in cfg.model_path.items(): - console.print(f"[dim]ONNX ({role}):[/dim] {path}") - elif cfg.model_path: - console.print(f"[dim]ONNX:[/dim] {cfg.model_path}") - if cfg.reference_path: - console.print(f"[dim]Reference:[/dim] {cfg.reference_path}") - console.print(f"[dim]Reference device:[/dim] {cfg.reference_device}") - if cfg.reference_ep: - console.print(f"[dim]Reference EP:[/dim] {cfg.reference_ep}") + console.print(f"[dim]Model ({role}):[/dim] {path}") + elif cfg.mode != "compare" and cfg.model_path: + console.print(f"[dim]Model:[/dim] {cfg.model_path}") # Metrics table console.print() diff --git a/src/winml/modelkit/commands/export.py b/src/winml/modelkit/commands/export.py index d575584e2..12cdb3021 100644 --- a/src/winml/modelkit/commands/export.py +++ b/src/winml/modelkit/commands/export.py @@ -26,11 +26,13 @@ import logging from pathlib import Path +from typing import TYPE_CHECKING import click from rich.console import Console from ..utils import cli as cli_utils +from ..utils.constants import EXPORT_TARGETS, ExportTarget from ..utils.logging import configure_logging from ..utils.model_input import ModelInputKind, classify_model_input @@ -39,6 +41,10 @@ console = Console() +if TYPE_CHECKING: + from ..export.cgc import CGCExporter + + def _delete_onnx_with_external_data(onnx_path: Path) -> None: """Delete an ONNX file and its external data files.""" import onnx @@ -90,8 +96,26 @@ def _warn_partial_composite(completed: list[Path]) -> None: required=True, help_text="HuggingFace model name or local path (e.g., prajjwal1/bert-tiny)", ) -@cli_utils.output_option("Output ONNX file path (e.g., model.onnx)", required=True) +@cli_utils.output_option("Final output path", required=True) @cli_utils.overwrite_option() +@click.option( + "--target", + type=click.Choice(list(EXPORT_TARGETS), case_sensitive=False), + default="onnx", + show_default=True, + help="Export target.", +) +@click.option( + "--options", + "target_options", + multiple=True, + metavar="KEY=VALUE", + help=( + "Additional target options (repeatable). CGIR supports external-weights, " + "update-opset, topo-sort-nodes, and freeze-dims " + "(for example freeze-dims=batch=1,seq=128)." + ), +) @click.option( "--batch-size", type=click.IntRange(min=1), @@ -171,6 +195,8 @@ def export( model: str, output: Path, overwrite: bool, + target: ExportTarget, + target_options: tuple[str, ...], batch_size: int, verbose: int, quiet: bool, @@ -239,17 +265,11 @@ def export( # Export only the encoder sub-model winml export -m google-t5/t5-small --task translation -o t5.onnx --submodel encoder """ - # Classify the -m value once (existence-first). Export only works with - # HuggingFace model IDs — reject ONNX files and folders early. + model_input = None if model: model_input = classify_model_input(model) if model_input.kind is ModelInputKind.INVALID: raise click.UsageError(model_input.error or f"Invalid model input: {model}") - if model_input.kind is ModelInputKind.ONNX_FILE: - raise click.UsageError( - "export requires a HuggingFace model ID, not an ONNX file. " - "Use 'winml inspect -m model.onnx' to inspect an existing ONNX model." - ) if model_input.kind is ModelInputKind.FOLDER: raise click.UsageError( "export requires a HuggingFace model ID, not a directory. " @@ -318,6 +338,33 @@ def export( export_config_dict = cli_utils.load_json_object(export_config, "--export-config") console.print(f"[dim]Loaded export config: {list(export_config_dict.keys())}[/dim]") + for settings in (_build_export_dict, export_config_dict): + if not cli_utils.is_cli_provided(ctx, "target") and "target" in settings: + target = settings["target"] + if not cli_utils.is_cli_provided(ctx, "target_options") and "options" in settings: + target_options = tuple(f"{key}={value}" for key, value in settings["options"].items()) + + exporter = _get_exporter(target, target_options) + + if model_input is not None and model_input.kind is ModelInputKind.ONNX_FILE: + if exporter is None: + raise click.UsageError( + "export requires a HuggingFace model ID, not an ONNX file. " + "Use 'winml inspect -m model.onnx' to inspect an existing ONNX model." + ) + try: + _guard_export_output(output, exporter, overwrite) + output.parent.mkdir(parents=True, exist_ok=True) + exporter.export_onnx( + model=Path(model_input.local_path or model), + output_path=output, + ) + return + except click.ClickException: + raise + except Exception as e: + raise click.ClickException(f"Export failed: {e}") from e + # Load shape overrides from JSON (task-independent). shape_overrides = None if shape_config: @@ -474,15 +521,27 @@ def _run_component_export(component_task: str | None, out_path: Path) -> None: else: console.print(f"[dim]Detected task: {detected_task}[/dim]") - export_stats = export_onnx( - model=pytorch_model, - output_path=out_path, - export_config=cfg, - model_id=model, - task=detected_task, - verbose=bool(verbose), - enable_reporting=with_report, - ) + export_stats: object + if exporter is None: + export_stats = export_onnx( + model=pytorch_model, + output_path=out_path, + export_config=cfg, + model_id=model, + task=detected_task, + verbose=bool(verbose), + enable_reporting=with_report, + ) + else: + export_stats = exporter.export_pytorch( + model=pytorch_model, + output_path=out_path, + export_config=cfg, + model_id=model, + task=detected_task, + verbose=bool(verbose), + enable_reporting=with_report, + ) logger.debug("Export stats: %s", export_stats) console.print(f"\n[bold green]Success![/bold green] Model exported to: {out_path}") @@ -500,8 +559,8 @@ def _run_component_export(component_task: str | None, out_path: Path) -> None: console.print(f" JSON: {json_metadata}") # Detect a composite pipeline (registry-driven). A composite fans out into one - # ONNX per sub-component, each written next to with a _ - # stem suffix; a plain model exports to the single output path as before. + # target artifact per sub-component, each written next to with a + # _ stem suffix; a plain model uses the single output path. # Detection suppresses only the expected "not a resolvable HF config" case # (OSError — e.g. the model reference isn't a hub id / has no local config); # intentional loud guards (empty registry, model-task incompatibility) and any @@ -550,7 +609,7 @@ def _run_component_export(component_task: str | None, out_path: Path) -> None: components = {submodel: components[submodel]} try: - console.print("\n[bold]Starting HTP export...[/bold]") + console.print("\n[bold]Starting export...[/bold]") if components: # A genuine multi-component fan-out can't take --input-specs (each @@ -580,7 +639,7 @@ def _run_component_export(component_task: str | None, out_path: Path) -> None: # Guard every target up front so an overwrite collision on a later # component can't leave an earlier one already written. for sub_out in sub_outputs.values(): - cli_utils.guard_output(sub_out, overwrite) + _guard_export_output(sub_out, exporter, overwrite) # Track sub-models this invocation actually completes. On a mid-run # failure we do NOT delete anything (the targets may be pre-existing @@ -600,7 +659,7 @@ def _run_component_export(component_task: str | None, out_path: Path) -> None: _warn_partial_composite(completed) raise else: - cli_utils.guard_output(output_path, overwrite) + _guard_export_output(output_path, exporter, overwrite) _run_component_export(task, output_path) except (click.UsageError, click.ClickException): @@ -613,3 +672,50 @@ def _run_component_export(component_task: str | None, out_path: Path) -> None: else: logger.error("Export failed: %s", e) raise click.ClickException(f"Export failed: {e}") from e + + +def _get_exporter( + target: ExportTarget, + target_options: tuple[str, ...], +) -> CGCExporter | None: + """Create the optional exporter for a non-default target.""" + from ..export.cgc import CGCExporter + + if target not in EXPORT_TARGETS: + raise click.UsageError(f"Invalid export target: {target!r}") + + exporter_types = { + "cgir": CGCExporter, + } + exporter_type = exporter_types.get(target) + if exporter_type is None: + if target_options: + raise click.UsageError("--options requires --target cgir") + return None + + try: + options = cli_utils.parse_options( + target_options, + exporter_type.options_type, + param_hint="--options", + ) + return exporter_type(options) + except click.ClickException: + raise + except Exception as e: + raise click.ClickException(f"Export failed: {e}") from e + + +def _guard_export_output( + output_path: Path, + exporter: CGCExporter | None, + overwrite: bool, +) -> None: + """Guard every artifact produced by the selected exporter.""" + artifacts = (output_path,) if exporter is None else exporter.output_artifacts(output_path) + for index, artifact in enumerate(artifacts): + cli_utils.guard_output( + artifact, + overwrite, + label="Output" if index == 0 else "Output sidecar", + ) diff --git a/src/winml/modelkit/commands/perf.py b/src/winml/modelkit/commands/perf.py index 8bbeccdea..bb3be9534 100644 --- a/src/winml/modelkit/commands/perf.py +++ b/src/winml/modelkit/commands/perf.py @@ -35,9 +35,13 @@ from ..utils.console import SafeConsole from ..utils.constants import ( ACCELERATOR_DEVICE_TYPES, + RUNTIME_BACKENDS, RUNTIME_NAMES, EPName, EPNameOrAlias, + RuntimeBackend, + RuntimeName, + resolve_runtime_api_backend, ) from ..utils.logging import ( configure_logging, @@ -62,7 +66,6 @@ from ..session.monitor.ep_monitor import WinMLEPMonitor from ..session.monitor.op_metrics import TraceFallbackReason from ..session.stats import PerfStats - from ..utils.constants import RuntimeName logger = logging.getLogger(__name__) @@ -77,13 +80,15 @@ def _resolve_runtime(runtime: RuntimeName, model: str) -> RuntimeName: - """Resolve ``auto`` from a local model folder, preserving explicit choices.""" + """Resolve ``auto`` from the model artifact, preserving explicit choices.""" if runtime != "auto": return runtime model_path = Path(model) if model_path.is_dir() and (model_path / "genai_config.json").is_file(): return "ort-genai" + if model_path.suffix.lower() == ".mlir": + return "winml-runtime" return "winml-ort" @@ -558,6 +563,8 @@ class BenchmarkConfig: """Configuration for benchmark execution.""" model_id: str + runtime: RuntimeName = "winml-ort" + backend: RuntimeBackend | None = None task: str | None = None submodel: str | None = None device: str = "auto" @@ -660,7 +667,10 @@ def to_dict(self) -> dict[str, Any]: result: dict[str, Any] = { "schema_version": 2, "benchmark_info": { - "runtime": _RUNTIME_TYPE, + "runtime": self.config.runtime, + "backend": resolve_runtime_api_backend( + self.config.runtime, self.config.model_id, self.config.backend + ), "model_id": self.config.model_id, "running_model_path": self.running_model_path, "task": self.actual_task, @@ -949,6 +959,9 @@ def __init__(self, config: BenchmarkConfig) -> None: self._effective_batch: int = config.batch_size self._memory: dict[str, float | None] | None = None self._memory_tracker: MemoryTracker | None = None + self._runtime_backend = resolve_runtime_api_backend( + config.runtime, config.model_id, config.backend + ) # Concrete device + EP resolved from the config's request, populated by # _resolve_device_ep() on the first call (before the build). The config # keeps the raw request (e.g. "auto"); these hold what actually drives @@ -967,6 +980,11 @@ def _resolve_device_ep(self) -> None: of them (WinMLAutoModel itself stays permissive: ep=None is a valid library mode). + For ``winml-runtime``, the Runtime native payload is loaded before + ``auto_device()`` can register an EP DLL. This pins the Runtime's matched + native dependencies and makes the load-order requirement part of target + resolution instead of relying on every caller to remember it. + Raises: ValueError: If the requested device/EP combination is unavailable or invalid (propagated from ``resolve_device``). @@ -985,7 +1003,8 @@ def _resolve_device_ep(self) -> None: ep=self.config.ep or "auto", device=self.config.device or "auto", source=self.config.ep_source, - ) + ), + backend=self._runtime_backend, ) self._ep_device = _resolve_perf_ep_device( target, self.config.device_luid, self.config.ep_options @@ -1177,7 +1196,11 @@ def _run_single(self) -> BenchmarkResult: } assert self._ep_device is not None pre_bench_kwargs = _pre_bench_kwargs_from_ep_device(self._ep_device, **pre_bench_common) - print_pre_bench_block(SafeConsole(stderr=True), **pre_bench_kwargs) + print_pre_bench_block( + SafeConsole(stderr=True), + runtime_api_backend=self._runtime_backend, + **pre_bench_kwargs, + ) # [3] Run benchmark if self.config.duration is not None: @@ -1216,19 +1239,19 @@ def _load_model(self) -> None: # Resolve the concrete device + EP first so a bad combo fails fast, # before from_pretrained/from_onnx kick off the build pipeline. - # This also binds ``self._ep_device`` via auto_device (loads the DLL). self._resolve_device_ep() assert self._ep_device is not None model_id = self.config.model_id model_path = Path(model_id) is_onnx = model_path.suffix.lower() == ".onnx" - if is_onnx and not model_path.exists(): + is_mlir = model_path.suffix.lower() == ".mlir" + if (is_onnx or is_mlir) and not model_path.exists(): # Surface a clear error for programmatic callers. The CLI guards # this earlier, but without this check from_pretrained would fall # through to HF loading and produce a confusing "not a valid JSON # file" error from AutoConfig. - raise FileNotFoundError(f"ONNX file not found: {model_path}") + raise FileNotFoundError(f"Model file not found: {model_path}") # Composite auto-detection. A bare seq2seq model such as T5 auto-detects # to a granular single-model task (text2text-generation) and would @@ -1245,7 +1268,7 @@ def _load_model(self) -> None: # bridges detection to that loadable pipeline task. Explicit --task and # ONNX inputs keep their resolved task untouched. resolved_task = self.config.task - if not is_onnx and resolved_task is None: + if not is_onnx and not is_mlir and resolved_task is None: from ..loader.resolution import resolve_composite_load_task try: @@ -1286,6 +1309,8 @@ def _load_model(self) -> None: "shape_config": self.config.shape_config, "allow_unsupported_nodes": self.config.allow_unsupported_nodes, "no_compile": self.config.no_compile, + "runtime": self.config.runtime, + "backend": self._runtime_backend, # optimize/analyze/max-optim toggles, forwarded by WinMLAutoModel to # build_hf_model / build_onnx_model. Shared mapping with build/eval. **cli_utils.build_pipeline_extra_kwargs( @@ -1303,6 +1328,15 @@ def _load_model(self) -> None: compile_provider_options=self.config.compile_ep_options, **common_kwargs, ) + elif is_mlir: + with suppress_native_warnings(enabled=True): + self._model = WinMLAutoModel.from_mlir( + mlir_path=model_path, + ep_device=self._ep_device, + task=resolved_task, + runtime=self.config.runtime, + backend="cgc", + ) else: with suppress_native_warnings(enabled=True): self._model = WinMLAutoModel.from_pretrained( @@ -2186,7 +2220,11 @@ def generate_output_path( under its own directory so per-sub-model reports don't collide. """ p = Path(model_id) - slug = p.stem if p.suffix.lower() == ".onnx" else model_id.replace("/", "_").replace("\\", "_") + slug = ( + p.stem + if p.suffix.lower() in {".onnx", ".mlir"} + else model_id.replace("/", "_").replace("\\", "_") + ) out_dir = Path.home() / ".cache" / "winml" / "perf" / slug if module_class: @@ -2746,9 +2784,18 @@ def _validate_duration( default="auto", show_default=True, help="'auto' selects ort-genai for folders containing genai_config.json, " - "otherwise winml-ort. 'winml-ort' benchmarks single-shot ONNX inference; " + "winml-runtime for .mlir files, otherwise winml-ort. " + "'winml-ort' benchmarks single-shot ONNX inference; " "'ort-genai' benchmarks an onnxruntime-genai bundle folder " - "(LLM generation: TTFT + decode tokens/sec).", + "(LLM generation: TTFT + decode tokens/sec); 'winml-runtime' performs " + "online conversion for ONNX/PyTorch inputs or loads CGC MLIR directly.", +) +@click.option( + "--backend", + type=click.Choice(list(RUNTIME_BACKENDS)), + default=None, + help="[winml-runtime] Execution backend for ONNX inputs (default: cgc). " + "MLIR inputs always use cgc.", ) @click.option( "--prompt", @@ -2958,6 +3005,7 @@ def perf( ctx: click.Context, model: str | None, runtime: RuntimeName, + backend: RuntimeBackend | None, prompt: str, prompt_file: Path | None, apply_template: bool, @@ -3068,6 +3116,10 @@ def perf( raise click.ClickException(f"Failed to resolve Hub-hosted ONNX path {model!r}: {e}") from e model = hf_model runtime = _resolve_runtime(runtime, model) + try: + effective_backend = resolve_runtime_api_backend(runtime, model, backend) + except ValueError as error: + raise click.UsageError(str(error)) from error # AC 11 (mockup spec): --top-k requires --op-tracing. Outside the # op-tracing section the flag is meaningless, so reject it explicitly # rather than silently ignoring a user's intent. @@ -3108,6 +3160,10 @@ def perf( elif "execution_provider" in cc: ep = (cc["execution_provider"], None) + if runtime == "winml-runtime" and ep_provider_options: + logger.warning("--ep-options are ignored with --runtime winml-runtime.") + ep_provider_options = None + json_mode = output_format == "json" console = SafeConsole(stderr=True) if json_mode else SafeConsole() @@ -3150,10 +3206,22 @@ def perf( # one source of truth. Rejects an invalid id up front; a path-shaped .onnx # that doesn't exist is caught below with a friendly "not found" message # (the pure classifier stays existence-agnostic). - model_input = classify_model_input(hf_model) - if model_input.kind is ModelInputKind.INVALID: - raise click.UsageError(model_input.error or f"Invalid model input: {hf_model}") - is_onnx = model_input.kind is ModelInputKind.ONNX_FILE + mlir_path = Path(hf_model) + is_mlir = mlir_path.suffix.lower() == ".mlir" + if is_mlir: + if runtime != "winml-runtime": + raise click.UsageError("MLIR inputs require --runtime winml-runtime.") + if not mlir_path.is_file(): + raise click.UsageError(f"MLIR file not found: {hf_model}") + if ep is not None: + logger.warning("--ep is ignored for MLIR inputs.") + ep = None + is_onnx = False + else: + model_input = classify_model_input(hf_model) + if model_input.kind is ModelInputKind.INVALID: + raise click.UsageError(model_input.error or f"Invalid model input: {hf_model}") + is_onnx = model_input.kind is ModelInputKind.ONNX_FILE if is_onnx and model_input.local_path and not Path(model_input.local_path).exists(): raise click.UsageError(f"ONNX file not found: {hf_model}") @@ -3410,6 +3478,8 @@ def perf( # ``ep_source_part`` were unpacked once from the --ep tuple above. config = BenchmarkConfig( model_id=hf_model, + runtime=runtime, + backend=effective_backend, task=task, submodel=submodel, device=device.lower(), diff --git a/src/winml/modelkit/compiler/configs.py b/src/winml/modelkit/compiler/configs.py index 2117b76d7..347ef46e9 100644 --- a/src/winml/modelkit/compiler/configs.py +++ b/src/winml/modelkit/compiler/configs.py @@ -182,6 +182,7 @@ def for_provider( "NvTensorRTRTXExecutionProvider": lambda: cls.for_nv_tensorrt_rtx(device=device), "OpenVINOExecutionProvider": lambda: cls.for_openvino(device=device), "VitisAIExecutionProvider": lambda: cls.for_vitisai(device=device), + "WinMLCGExecutionProvider": lambda: cls.for_winmlcg(device=device), "MIGraphXExecutionProvider": cls.for_migraphx, "CPUExecutionProvider": cls.for_cpu, } @@ -263,6 +264,17 @@ def for_openvino(cls, device: str | None = None) -> WinMLCompileConfig: ) return cls(ep_config=ep_cfg) + @classmethod + def for_winmlcg(cls, device: str | None = None) -> WinMLCompileConfig: + """Factory for Windows ML Compute Graph EP compilation.""" + return cls( + ep_config=EPConfig( + provider="winmlcg", + enable_ep_context=True, + device=device or "gpu", + ) + ) + @classmethod def for_vitisai(cls, device: str | None = None) -> WinMLCompileConfig: """Factory for Vitis AI (AMD NPU) compilation. diff --git a/src/winml/modelkit/config/build.py b/src/winml/modelkit/config/build.py index 020c58033..6ed6d3478 100644 --- a/src/winml/modelkit/config/build.py +++ b/src/winml/modelkit/config/build.py @@ -63,6 +63,7 @@ from ..optim.config import WinMLOptimizationConfig from ..quant.config import WinMLQuantizationConfig from ..utils.config_utils import merge_config +from ..utils.constants import normalize_ep_name # NOTE: WinMLEvaluationConfig is imported lazily to avoid pulling @@ -78,7 +79,7 @@ from torch import nn from ..eval.config import WinMLEvaluationConfig # noqa: TC004 - from ..utils.constants import EPNameOrAlias + from ..utils.constants import EPNameOrAlias, RuntimeBackend ExportPolicyTargetRequest = tuple[str | None, str | None] @@ -108,7 +109,9 @@ class WinMLBuildConfig: optim: Optimization configuration quant: Quantization configuration compile: Compilation configuration + convert: Optional ONNX-to-MLIR export configuration for the build CLI eval: Evaluation configuration + is_cgc: Whether the build uses CGC compatibility optimization policy Example: from winml.modelkit.config import WinMLBuildConfig @@ -143,6 +146,8 @@ class WinMLBuildConfig: auto: bool = True # Skip ORT optimization. Pre-quantized inputs also clear ``quant``. skip_optimize: bool = False + convert: WinMLExportConfig | None = None + is_cgc: bool = False def __post_init__(self) -> None: # Lazy import: inject into module globals so typing.get_type_hints() @@ -160,6 +165,7 @@ def from_dict(cls, config_dict: dict) -> WinMLBuildConfig: export_data = config_dict.get("export", {}) quant_data = config_dict.get("quant") compile_data = config_dict.get("compile") + convert_data = config_dict.get("convert") eval_data = config_dict.get("eval") eval_cfg = None if eval_data is not None: @@ -177,6 +183,10 @@ def from_dict(cls, config_dict: dict) -> WinMLBuildConfig: eval=eval_cfg, auto=config_dict.get("auto", True), skip_optimize=config_dict.get("skip_optimize", False), + convert=( + WinMLExportConfig.from_dict(convert_data) if convert_data is not None else None + ), + is_cgc=config_dict.get("is_cgc", False), ) def to_dict(self) -> dict: @@ -186,6 +196,8 @@ def to_dict(self) -> dict: result["auto"] = False if self.skip_optimize: result["skip_optimize"] = True + if self.is_cgc: + result["is_cgc"] = True result.update( { "export": self.export.to_dict() if self.export is not None else None, @@ -200,6 +212,8 @@ def to_dict(self) -> dict: result["loader"] = loader_dict if self.eval is not None: result["eval"] = self.eval.to_dict() + if self.convert is not None: + result["convert"] = self.convert.to_dict() return result def validate(self) -> None: @@ -251,6 +265,9 @@ def validate(self) -> None: ): errors.append("compile.ep_config.provider is required when compile is enabled") + if self.convert is not None and self.convert.target != "cgir": + errors.append("convert.target must be 'cgir'") + if errors: raise ValueError("Invalid WinMLBuildConfig:\n" + "\n".join(f" - {e}" for e in errors)) @@ -385,6 +402,7 @@ def _apply_target_policy( device: str, precision: str, ep: str | None, + backend: RuntimeBackend | None = None, ) -> None: """Apply resolved device/precision policy to quant and compile sections.""" from ..sysinfo.hardware import get_available_devices @@ -409,6 +427,7 @@ def _apply_target_policy( ep=resolved_ep, available_devices=available_devices, task=config.loader.task, + backend=backend, ) # Mutate quant in place so calibration identity fields stamped by @@ -488,6 +507,7 @@ def resolve_quant_compile_config( precision: str = "auto", ep: str | None = None, task: str | None = None, + backend: RuntimeBackend | None = None, ) -> tuple[WinMLQuantizationConfig | None, WinMLCompileConfig | None]: """Resolve quantization and compilation config from device/precision policy. @@ -501,6 +521,7 @@ def resolve_quant_compile_config( "int16", or "w{x}a{y}" e.g. "w8a16"). ep: Explicit execution provider override. task: Model task (used for precision heuristics, e.g., LLM on GPU). + backend: Runtime backend used to resolve default precision and compilation. Returns: Tuple of (quant_config, compile_config). Either may be None when the @@ -528,6 +549,7 @@ def resolve_quant_compile_config( ep=resolved_ep, available_devices=available_devices, task=task, + backend=backend, ) if policy.device == "auto": @@ -565,6 +587,20 @@ def resolve_quant_compile_config( # ============================================================================= +def _apply_cgc_config( + config: WinMLBuildConfig, *, backend: RuntimeBackend | None, ep: str | None +) -> None: + """Populate CGC optimization and conversion stages and disable compilation.""" + if backend == "cgc" or normalize_ep_name(ep) == "WinMLCGExecutionProvider": + config.is_cgc = True + config.skip_optimize = False + config.auto = False + config.optim = WinMLOptimizationConfig.for_cgc() + config.compile = None + if backend == "cgc" and config.convert is None: + config.convert = WinMLExportConfig(target="cgir") + + def generate_onnx_build_config( onnx_path: str | Path, *, @@ -574,6 +610,7 @@ def generate_onnx_build_config( ep: str | None = None, override: BuildConfigOverride | None = None, no_compile: bool = False, + backend: RuntimeBackend | None = None, ) -> WinMLBuildConfig: """Generate build config for a pre-exported ONNX model (Scenario D). @@ -625,6 +662,7 @@ def generate_onnx_build_config( precision=precision, ep=ep, task=task, + backend=backend, ) if is_quantized_onnx(onnx_path_resolved): @@ -647,6 +685,8 @@ def generate_onnx_build_config( # "already exported, skip export stage". config.export = None + _apply_cgc_config(config, backend=backend, ep=ep) + # no_compile overrides policy and override — applied last so it always wins if no_compile: config.compile = None @@ -828,6 +868,7 @@ def generate_hf_build_config( export_policy_target: ExportPolicyTargetRequest | None = None, policy_overrides_config: bool = False, no_compile: bool = False, + backend: RuntimeBackend | None = None, ) -> WinMLBuildConfig: ... @@ -849,6 +890,7 @@ def generate_hf_build_config( export_policy_target: ExportPolicyTargetRequest | None = None, policy_overrides_config: bool = False, no_compile: bool = False, + backend: RuntimeBackend | None = None, ) -> list[WinMLBuildConfig]: ... @@ -874,6 +916,7 @@ def generate_hf_build_config( export_policy_target: ExportPolicyTargetRequest | None = None, policy_overrides_config: bool = False, no_compile: bool = False, + backend: RuntimeBackend | None = None, ) -> WinMLBuildConfig | list[WinMLBuildConfig]: ... @@ -894,6 +937,7 @@ def generate_hf_build_config( export_policy_target: ExportPolicyTargetRequest | None = None, policy_overrides_config: bool = False, no_compile: bool = False, + backend: RuntimeBackend | None = None, ) -> WinMLBuildConfig | list[WinMLBuildConfig]: """Generate WinMLBuildConfig for a HuggingFace model (Scenarios A/B/C). @@ -1059,6 +1103,7 @@ class name. Uses torchinfo to discover submodules and infer device=device, precision=precision, ep=ep, + backend=backend, ) if override: @@ -1102,8 +1147,11 @@ class name. Uses torchinfo to discover submodules and infer device=device, precision=precision, ep=ep, + backend=backend, ) + _apply_cgc_config(parent_config, backend=backend, ep=ep) + # no_compile overrides policy — applied last so it always wins if no_compile: parent_config.compile = None @@ -1148,7 +1196,10 @@ class name. Uses torchinfo to discover submodules and infer ) logger.info("Found %d submodules matching '%s'", len(submodules), module) - return [_build_submodule_config(sub_info, parent_config) for sub_info in submodules] + return [ + _build_submodule_config(sub_info, parent_config, backend=backend, ep=ep) + for sub_info in submodules + ] return parent_config @@ -1175,6 +1226,7 @@ def generate_build_config( ep: str | None = None, export_policy_target: ExportPolicyTargetRequest | None = None, onnx_path: str | Path | None = None, + backend: RuntimeBackend | None = None, ) -> WinMLBuildConfig: ... @@ -1195,6 +1247,7 @@ def generate_build_config( ep: str | None = None, export_policy_target: ExportPolicyTargetRequest | None = None, onnx_path: str | Path | None = None, + backend: RuntimeBackend | None = None, ) -> list[WinMLBuildConfig]: ... @@ -1214,6 +1267,7 @@ def generate_build_config( ep: str | None = None, export_policy_target: ExportPolicyTargetRequest | None = None, onnx_path: str | Path | None = None, + backend: RuntimeBackend | None = None, ) -> WinMLBuildConfig | list[WinMLBuildConfig]: """Generate WinMLBuildConfig by orchestrating existing modules. @@ -1253,6 +1307,7 @@ class name (HF path only). precision=precision, ep=ep, override=override, + backend=backend, ) # Single call resolves against generate_hf_build_config's `module: str | None` # overload, which returns WinMLBuildConfig | list[WinMLBuildConfig] — matching @@ -1273,6 +1328,7 @@ class name (HF path only). ep=ep, export_policy_target=export_policy_target, policy_overrides_config=True, + backend=backend, ) @@ -1284,6 +1340,9 @@ class name (HF path only). def _build_submodule_config( sub_info: SubmoduleInfo, parent_config: WinMLBuildConfig, + *, + backend: RuntimeBackend | None = None, + ep: str | None = None, ) -> WinMLBuildConfig: """Build a WinMLBuildConfig for a single discovered submodule. @@ -1330,7 +1389,7 @@ def _input_name(i: int) -> str: OutputTensorSpec(name=f"output_{i}") for i in range(len(sub_info.output_shapes)) ] - return WinMLBuildConfig( + config = WinMLBuildConfig( loader=WinMLLoaderConfig( # task intentionally omitted — submodules don't have tasks model_type=parent_config.loader.model_type, @@ -1368,6 +1427,8 @@ def _input_name(i: int) -> str: ), compile=copy.deepcopy(parent_config.compile), ) + _apply_cgc_config(config, backend=backend, ep=ep) + return config def _merge_export_config( diff --git a/src/winml/modelkit/config/precision.py b/src/winml/modelkit/config/precision.py index 57cb26b61..a4e33dc5c 100644 --- a/src/winml/modelkit/config/precision.py +++ b/src/winml/modelkit/config/precision.py @@ -33,7 +33,7 @@ if TYPE_CHECKING: # Referenced only from the quoted ``cast()`` below, so importing it at # runtime would leave an unused import behind. - from ..utils.constants import EPNameOrAlias + from ..utils.constants import EPNameOrAlias, RuntimeBackend logger = logging.getLogger(__name__) @@ -346,6 +346,7 @@ def resolve_precision( ep: str | None = None, available_devices: list[str] | None = None, task: str | None = None, + backend: RuntimeBackend | None = None, ) -> PrecisionPolicy: """Resolve precision into a concrete PrecisionPolicy. @@ -367,6 +368,8 @@ def resolve_precision( available_devices: Prioritized device list from sysinfo.get_available_devices(). Used when device="auto" + precision is explicit. task: Optional task name for LLM-specific warnings. + backend: CGC defaults auto precision to FP16 and disables offline compilation. + Explicit precision choices are preserved. Returns: PrecisionPolicy with all fields resolved. @@ -403,6 +406,10 @@ def resolve_precision( elif device not in supported_devices: raise ValueError(f"EP '{ep}' does not support device '{device}'.") + is_cgc = backend == "cgc" or ep == "WinMLCGExecutionProvider" + if is_cgc and resolved_precision == "auto": + resolved_precision = "fp16" + # --- Both auto: no-op, keep config defaults --- if device == "auto" and resolved_precision == "auto": return PrecisionPolicy( @@ -469,7 +476,9 @@ def resolve_precision( # The policy contract uses short aliases, with CPU represented as no # offline compiler. - compile_provider = ep_short_or_none(effective_ep) if effective_ep is not None else None + compile_provider = ( + ep_short_or_none(effective_ep) if effective_ep is not None and not is_cgc else None + ) # Resolve weight/activation types — supports named presets and w{x}a{y}. # Weight-only precisions (int4, w4a16) use RTN, not QDQ — they have no diff --git a/src/winml/modelkit/datasets/random_dataset.py b/src/winml/modelkit/datasets/random_dataset.py index 105ca8c40..dba988b4f 100644 --- a/src/winml/modelkit/datasets/random_dataset.py +++ b/src/winml/modelkit/datasets/random_dataset.py @@ -32,7 +32,7 @@ class RandomDataset: model_path: Path to ONNX model file max_samples: Maximum number of samples to generate (default: 100) seed: Random seed for reproducible data generation (default: 42) - **kwargs: Additional keyword arguments (ignored) + **kwargs: Additional keyword arguments """ TASK_TYPE = "random" @@ -40,7 +40,7 @@ class RandomDataset: def __init__( self, - model_path: str, + model_path: str | None, max_samples: int = 100, seed: int = 42, **kwargs: Any, @@ -52,7 +52,9 @@ def __init__( # Cache io_config (loads ONNX once) from ..onnx import get_io_config - self._io_config = get_io_config(model_path) + self._io_config = ( + get_io_config(model_path) if model_path is not None else kwargs["io_config"] + ) # Build InputTensorSpec list for reuse across samples from ..onnx import InputTensorSpec diff --git a/src/winml/modelkit/ep_path.py b/src/winml/modelkit/ep_path.py index df70982c4..bec8bcd35 100644 --- a/src/winml/modelkit/ep_path.py +++ b/src/winml/modelkit/ep_path.py @@ -521,6 +521,9 @@ def resolve(self) -> Iterator[EPEntry]: path = Path(str(dist.locate_file(rel))) if not path.exists(): + # Temporary workaround: remove after WinMLCG EP is released. + if path.name == "WinMLCGEp.dll": + return logger.warning( "PyPISource: distribution %r installed but DLL missing at %s", self.distribution, @@ -1575,6 +1578,11 @@ def _default_ep_sources() -> list[EPSource]: """ return [ # 1. Manually installed PyPI plugin wheels — BYO sources. + PyPISource( + distribution="windowsml", + relative_dll="windowsml/lib/WinMLCGEp.dll", + eps=("WinMLCGExecutionProvider",), + ), PyPISource( distribution="onnxruntime-ep-openvino", relative_dll=("onnxruntime_ep_openvino/onnxruntime_providers_openvino_plugin.dll"), diff --git a/src/winml/modelkit/eval/config.py b/src/winml/modelkit/eval/config.py index 2cb4f3dcc..d4c6073ee 100644 --- a/src/winml/modelkit/eval/config.py +++ b/src/winml/modelkit/eval/config.py @@ -11,11 +11,11 @@ from pathlib import Path from typing import Any, Literal -from ..utils.constants import EPNameOrAlias +from ..utils.constants import EPNameOrAlias, RuntimeBackend from ..utils.eval_utils import EvalMode -EvalRuntime = Literal["winml-ort", "pytorch"] +EvalRuntime = Literal["winml-ort", "winml-runtime", "pytorch"] @dataclass @@ -107,7 +107,7 @@ class WinMLEvaluationConfig: Attributes: model_id: HuggingFace model ID for config/preprocessor resolution. - model_path: Path to .onnx model file, or a ``{role: path}`` dict for + model_path: Path to an ONNX or MLIR model file, or a ``{role: path}`` dict for composite models (e.g. ``{"image-encoder": "...", "text-encoder": "..."}``). None = build from model_id. input_data: Path to a ``.npz`` archive of real input tensors for @@ -116,7 +116,7 @@ class WinMLEvaluationConfig: randomly generated ones. The leading axis of each array is the sample axis, so one archive can hold ``N`` samples; all inputs must share the same leading length. - reference_path: Path to a second ``.onnx`` file used as the reference in + reference_path: Path to an ONNX file used as the reference in ``--mode compare``. When set, both ``model_path`` and ``reference_path`` load as WinML model instances and their output tensors are compared directly, so no ``model_id`` / ``task`` / HF reference is needed. @@ -143,12 +143,13 @@ class WinMLEvaluationConfig: - ``"winml-ort"`` (default): export Hugging Face checkpoints to ONNX and evaluate with WinML. + - ``"winml-runtime"``: evaluate a pre-built CGC MLIR artifact. - ``"pytorch"``: evaluate the original Hugging Face checkpoint. mode: Evaluation mode (see :data:`EvalMode`). - - ``"onnx"`` (default): evaluate the ONNX candidate on the + - ``"onnx"`` (default): evaluate the candidate model on the labeled dataset. - - ``"compare"``: compare ONNX vs HF reference output tensors + - ``"compare"``: compare candidate vs reference output tensors on identical random inputs and report tensor-similarity metrics per output tensor. When ``reference_path`` is set, the reference is a second ONNX file instead of the HF model. @@ -193,6 +194,7 @@ class WinMLEvaluationConfig: use_cache: bool = True rebuild: bool = False runtime: EvalRuntime = "winml-ort" + backend: RuntimeBackend | None = None trust_remote_code: bool = False _auto_device_selected: bool = field(default=False, repr=False, compare=False, kw_only=True) _pipeline_device_override: str | None = field( @@ -214,6 +216,8 @@ def pipeline_device(self) -> str: def to_dict(self) -> dict: """Convert to dictionary for serialization.""" result: dict = {"runtime": self.runtime} + if self.backend is not None: + result["backend"] = self.backend if self.model_id is not None: result["model_id"] = self.model_id if self.model_path is not None: @@ -257,7 +261,7 @@ def to_dict(self) -> dict: result["output_path"] = str(self.output_path) if self.mode != "onnx": result["mode"] = self.mode - if self.runtime == "winml-ort": + if self.runtime in ("winml-ort", "winml-runtime"): result["skip_build"] = self.skip_build result["use_cache"] = self.use_cache result["rebuild"] = self.rebuild @@ -309,5 +313,6 @@ def from_dict(cls, data: dict) -> WinMLEvaluationConfig: use_cache=data.get("use_cache", True), rebuild=data.get("rebuild", False), runtime=data.get("runtime", "winml-ort"), + backend=data.get("backend"), trust_remote_code=data.get("trust_remote_code", False), ) diff --git a/src/winml/modelkit/eval/evaluate.py b/src/winml/modelkit/eval/evaluate.py index 6e5b4f223..2226eaa6f 100644 --- a/src/winml/modelkit/eval/evaluate.py +++ b/src/winml/modelkit/eval/evaluate.py @@ -16,6 +16,7 @@ from rich.console import Console +from ..utils.constants import resolve_runtime_api_backend from .config import WinMLEvaluationConfig @@ -37,6 +38,7 @@ class _ModelLoaderKind(Enum): PYTORCH = auto() GENAI = auto() EVALUATOR_MANAGED = auto() + MLIR = auto() ONNX = auto() PRETRAINED = auto() @@ -45,6 +47,11 @@ def _select_model_loader(config: WinMLEvaluationConfig) -> _ModelLoaderKind: """Select the model-loading path shared by loading and CLI diagnostics.""" if config.runtime == "pytorch": return _ModelLoaderKind.PYTORCH + if ( + isinstance(config.model_path, str) + and config.model_path.lower().endswith(".mlir") + ): + return _ModelLoaderKind.MLIR if config.task == "text-generation": return _ModelLoaderKind.GENAI if isinstance(config.model_path, dict) and config.task == "mask-generation": @@ -120,7 +127,7 @@ def get_evaluator_class(config: WinMLEvaluationConfig) -> type[WinMLEvaluator]: def _validate_pytorch_runtime_config(config: WinMLEvaluationConfig) -> None: """Validate state that cannot apply to the PyTorch runtime.""" - if config.runtime == "winml-ort": + if config.runtime != "pytorch": return incompatible: list[str] = [] @@ -381,18 +388,29 @@ def load_model( ) model_id = config.model_id - if model_id is None and loader is not _ModelLoaderKind.ONNX: + if model_id is None and loader not in ( + _ModelLoaderKind.MLIR, + _ModelLoaderKind.ONNX, + ): raise ValueError("model_id is required.") if loader is _ModelLoaderKind.EVALUATOR_MANAGED: # Evaluator-driven session loading; skip WinMLAutoModel entirely. return None + assert config.runtime != "pytorch" + # Resolve EPDeviceTarget then bind a WinMLEPDevice at the boundary. Eval # config carries an optional ep field; resolve_device deduces device/ep # when either is 'auto'. device = (config.device or "auto").lower() - target = resolve_device(EPDeviceTarget(ep=config.ep or "auto", device=device)) + runtime_backend = resolve_runtime_api_backend( + config.runtime, config.model_path, config.backend + ) + target = resolve_device( + EPDeviceTarget(ep=config.ep or "auto", device=device), + backend=runtime_backend, + ) registry = WinMLEPRegistry.instance() ep_device = ( registry.auto_device(target, device_luid=config.device_luid) @@ -403,6 +421,28 @@ def load_model( from onnxruntime.capi.onnxruntime_pybind11_state import RuntimeException try: + if loader is _ModelLoaderKind.MLIR: + hf_config = None + if config.model_id is not None: + from transformers import AutoConfig + + from ..loader import load_hf_config + + hf_config = load_hf_config( + AutoConfig, + config.model_id, + trust_remote_code=config.trust_remote_code, + ) + mlir_model = WinMLAutoModel.from_mlir( + mlir_path=cast("str", config.model_path), + ep_device=ep_device, + task=config.task, + runtime="winml-runtime", + backend="cgc", + ) + mlir_model.config = hf_config + return mlir_model + if loader is _ModelLoaderKind.ONNX: # Pre-built ONNX: precision is already baked into the model and is # ignored here (mirrors winml perf's ONNX path). @@ -419,20 +459,22 @@ def load_model( if model_id is not None else None ) - model = WinMLAutoModel.from_onnx( + onnx_model = WinMLAutoModel.from_onnx( # ``model_path`` is narrowed to ``str | dict[str, str]`` here; # cast bridges dict value-type invariance (str vs str | Path). onnx_path=cast("str | dict[str, str | Path]", config.model_path), ep_device=ep_device, task=config.task, + runtime=config.runtime, + backend=runtime_backend, skip_build=config.skip_build, config=quant_override, hf_config=hf_config, **cache_kwargs, **pipeline_kwargs, ) - model.config = hf_config - return model + onnx_model.config = hf_config + return onnx_model assert model_id is not None @@ -456,6 +498,8 @@ def load_model( task=config.task, device=config.device, ep=config.ep, + runtime=config.runtime, + backend=runtime_backend, precision=config.precision, allow_unsupported_nodes=config.allow_unsupported_nodes, config=build_override, @@ -660,8 +704,17 @@ def evaluate( """ from ..utils.eval_utils import EVAL_MODES - if config.runtime not in ("winml-ort", "pytorch"): - raise ValueError(f"Invalid runtime {config.runtime!r}; expected 'winml-ort' or 'pytorch'.") + if config.runtime not in ("winml-ort", "winml-runtime", "pytorch"): + raise ValueError( + f"Invalid runtime {config.runtime!r}; expected 'winml-ort', " + "'winml-runtime', or 'pytorch'." + ) + is_mlir = ( + isinstance(config.model_path, str) + and config.model_path.lower().endswith(".mlir") + ) + if is_mlir and config.runtime != "winml-runtime": + raise ValueError("MLIR inputs require runtime='winml-runtime'.") if pytorch_model is not None: config = _prepare_supplied_pytorch_model(config, pytorch_model) mode = config.mode if config.mode is not None else "onnx" @@ -755,21 +808,46 @@ def evaluate( def print_config(config: WinMLEvaluationConfig) -> None: """Print effective evaluation config to the console (quantize.py style).""" ds = config.dataset + backend = resolve_runtime_api_backend( + config.runtime, config.model_path, config.backend + ) output_console = Console() - if config.model_id is not None: + if config.mode != "compare" and config.model_id is not None: output_console.print(f"[bold blue]Model:[/bold blue] {config.model_id}") - if config.model_path is not None: + if config.mode != "compare" and config.model_path is not None: output_console.print(f"[bold blue]Model path:[/bold blue] {config.model_path}") if config.input_data is not None: output_console.print(f"[bold blue]Input data:[/bold blue] {config.input_data}") - if config.reference_path is not None: - output_console.print(f"[bold blue]Reference:[/bold blue] {config.reference_path}") if config.task is not None: output_console.print(f"[bold blue]Task:[/bold blue] {config.task}") - output_console.print(f"[bold blue]Runtime:[/bold blue] {config.runtime}") - output_console.print(f"[bold blue]Device:[/bold blue] {config.device}") - if config.ep is not None: - output_console.print(f"[bold blue]EP:[/bold blue] {config.ep}") + if config.mode == "compare": + output_console.print( + f"[bold blue]Candidate:[/bold blue] {config.model_path or config.model_id}" + ) + output_console.print(f"[bold blue]Candidate runtime:[/bold blue] {config.runtime}") + output_console.print(f"[bold blue]Candidate device:[/bold blue] {config.device}") + if backend != "cgc": + output_console.print(f"[bold blue]Candidate EP:[/bold blue] {config.ep or 'auto'}") + output_console.print( + f"[bold blue]Reference:[/bold blue] {config.reference_path or config.model_id}" + ) + output_console.print( + f"[bold blue]Reference runtime:[/bold blue] " + f"{'winml-ort' if config.reference_path else 'pytorch'}" + ) + output_console.print( + f"[bold blue]Reference device:[/bold blue] " + f"{config.reference_device if config.reference_path else 'cpu'}" + ) + output_console.print( + f"[bold blue]Reference EP:[/bold blue] " + f"{(config.reference_ep or 'auto') if config.reference_path else 'n/a'}" + ) + else: + output_console.print(f"[bold blue]Runtime:[/bold blue] {config.runtime}") + output_console.print(f"[bold blue]Device:[/bold blue] {config.device}") + if backend != "cgc" and config.ep is not None: + output_console.print(f"[bold blue]EP:[/bold blue] {config.ep}") if config.runtime == "winml-ort": output_console.print(f"[bold blue]Precision:[/bold blue] {config.precision}") if config.mode != "compare": diff --git a/src/winml/modelkit/eval/metrics/__init__.py b/src/winml/modelkit/eval/metrics/__init__.py index 0488db527..e7b340413 100644 --- a/src/winml/modelkit/eval/metrics/__init__.py +++ b/src/winml/modelkit/eval/metrics/__init__.py @@ -21,6 +21,7 @@ from .mean_iou import IGNORE_INDEX, MeanIoUMetric from .pseudo_perplexity import PseudoPerplexityMetric from .spearman_correlation import SpearmanCorrelationMetric + from .tensor_similarity import TensorSimilarityMetric from .top_k_accuracy import TopKAccuracyMetric @@ -38,6 +39,7 @@ "MeanIoUMetric": ".mean_iou:MeanIoUMetric", "PseudoPerplexityMetric": ".pseudo_perplexity:PseudoPerplexityMetric", "SpearmanCorrelationMetric": ".spearman_correlation:SpearmanCorrelationMetric", + "TensorSimilarityMetric": ".tensor_similarity:TensorSimilarityMetric", "TopKAccuracyMetric": ".top_k_accuracy:TopKAccuracyMetric", } @@ -69,5 +71,6 @@ def __dir__() -> list[str]: "MeanIoUMetric", "PseudoPerplexityMetric", "SpearmanCorrelationMetric", + "TensorSimilarityMetric", "TopKAccuracyMetric", ] diff --git a/src/winml/modelkit/eval/tensor_similarity_evaluator.py b/src/winml/modelkit/eval/tensor_similarity_evaluator.py index 678afd5cf..89a696af2 100644 --- a/src/winml/modelkit/eval/tensor_similarity_evaluator.py +++ b/src/winml/modelkit/eval/tensor_similarity_evaluator.py @@ -45,6 +45,7 @@ def _make_reference_config(config: WinMLEvaluationConfig) -> WinMLEvaluationConf model_path=config.reference_path, reference_path=None, runtime="winml-ort", + backend=None, device=config.reference_device, device_luid=config.reference_device_luid, ep=config.reference_ep, @@ -61,6 +62,7 @@ def _make_reference_config(config: WinMLEvaluationConfig) -> WinMLEvaluationConf model_path=None, reference_path=None, runtime="pytorch", + backend=None, device="cpu", device_luid=None, ep=None, @@ -123,7 +125,8 @@ def prepare_data(self) -> Any: ds = self.config.dataset return RandomDataset( - model_path=str(self.model.onnx_path), + model_path=None, + io_config=self.model.io_config, max_samples=int(ds.samples if ds.samples is not None else 100), seed=int(ds.seed if ds.seed is not None else 42), ) diff --git a/src/winml/modelkit/export/cgc/__init__.py b/src/winml/modelkit/export/cgc/__init__.py new file mode 100644 index 000000000..76eb0682f --- /dev/null +++ b/src/winml/modelkit/export/cgc/__init__.py @@ -0,0 +1,10 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""CGC export backend.""" + +from .exporter import CGCExporter, CGCExportResult, CGCOptions, export_cgc + + +__all__ = ["CGCExportResult", "CGCExporter", "CGCOptions", "export_cgc"] diff --git a/src/winml/modelkit/export/cgc/artifacts.py b/src/winml/modelkit/export/cgc/artifacts.py new file mode 100644 index 000000000..46973f0af --- /dev/null +++ b/src/winml/modelkit/export/cgc/artifacts.py @@ -0,0 +1,12 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""CGC artifact path conventions.""" + +from pathlib import Path + + +def cgc_metadata_path(model_path: Path) -> Path: + """Return the CLI metadata sidecar path for a CGC model.""" + return model_path.with_name(f"{model_path.stem}_metadata.json") diff --git a/src/winml/modelkit/export/cgc/exporter.py b/src/winml/modelkit/export/cgc/exporter.py new file mode 100644 index 000000000..6518517f9 --- /dev/null +++ b/src/winml/modelkit/export/cgc/exporter.py @@ -0,0 +1,373 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""ONNX-to-CGC export orchestration.""" + +from __future__ import annotations + +import json +import tempfile +import time +from dataclasses import dataclass, replace +from pathlib import Path +from typing import TYPE_CHECKING, Any + +import onnx +from rich.console import Console + +from .artifacts import cgc_metadata_path +from .foundry import FoundryCompileError, FoundryCompiler + + +if TYPE_CHECKING: + from torch import nn + + from ..config import WinMLExportConfig + + +@dataclass(frozen=True) +class CGCOptions: + """Configuration shared by CGC export steps.""" + + external_weights: bool = False + topo_sort_nodes: bool = True + update_opset: bool = True + freeze_dims: str = "" + + +@dataclass(frozen=True) +class CGCExportResult: + """Metadata produced by a CGC export.""" + + input_names: tuple[str, ...] + output_names: tuple[str, ...] + export_stats: dict[str, Any] | None = None + + +class CGCExporter: + """Generate standalone CGC Input IR.""" + + options_type = CGCOptions + + def __init__(self, options: CGCOptions) -> None: + """Initialize the exporter with typed CGC options.""" + self.options = options + self._freeze_dims = _parse_freeze_dims(options.freeze_dims) + self.console = Console(width=100, highlight=False) + + def export_pytorch( + self, + *, + model: nn.Module, + output_path: str | Path, + export_config: WinMLExportConfig, + model_id: str, + task: str | None, + verbose: bool, + enable_reporting: bool, + **onnx_kwargs: Any, + ) -> CGCExportResult: + """Export PyTorch to CGC without exposing the intermediate representation.""" + from .. import export_pytorch as export_onnx + + output_path = Path(output_path) + output_path.parent.mkdir(parents=True, exist_ok=True) + with tempfile.TemporaryDirectory( + prefix=".winml-cgc-source-", + dir=output_path.parent, + ) as temporary: + intermediate_path = Path(temporary) / "model.onnx" + export_stats = export_onnx( + model=model, + output_path=intermediate_path, + export_config=export_config, + model_id=model_id, + task=task, + verbose=verbose, + enable_reporting=enable_reporting, + **onnx_kwargs, + ) + result = self.export_onnx( + model=intermediate_path, + output_path=output_path, + ) + if enable_reporting: + self._publish_onnx_reports(intermediate_path, output_path) + return replace(result, export_stats=export_stats) + + def export_onnx( + self, + model: str | Path, + output_path: str | Path, + ) -> CGCExportResult: + """Export an ONNX model to standalone CGC MLIR.""" + source_path = Path(model) + if not source_path.is_file(): + raise FileNotFoundError(f"ONNX model not found: {source_path}") + + output_path = Path(output_path) + output_path.parent.mkdir(parents=True, exist_ok=True) + + self.console.print("\n" + "=" * 80) + self.console.print("🚀 [bold cyan]ONNX TO CGC EXPORT PROCESS[/bold cyan]") + self.console.print("=" * 80) + self.console.print(f"[bold blue]Input:[/bold blue] {source_path}") + self.console.print(f"[bold blue]Output:[/bold blue] {output_path}") + self.console.print("[bold blue]Format:[/bold blue] CGC MLIR") + export_start = time.monotonic() + + model_proto = onnx.load(str(source_path), load_external_data=False) + initializer_names = { + initializer.name for initializer in model_proto.graph.initializer + } + initializer_names.update( + initializer.values.name + for initializer in model_proto.graph.sparse_initializer + ) + self._freeze_dims = _parse_freeze_dims(self.options.freeze_dims) + self._freeze_batch_size(model_proto, initializer_names) + result = CGCExportResult( + input_names=tuple( + value.name + for value in model_proto.graph.input + if value.name not in initializer_names + ), + output_names=tuple(value.name for value in model_proto.graph.output), + ) + + self._export_mlir( + source_path, + output_path, + result, + ) + + self.console.print("\n[bold green]✅ CGC EXPORT COMPLETE[/bold green]") + self.console.print(f"[dim]Total time: {time.monotonic() - export_start:.2f}s[/dim]") + return result + + def _freeze_batch_size( + self, model: onnx.ModelProto, initializer_names: set[str], + ) -> None: + """Supplement explicit dimension overrides with a detected batch_size default.""" + if "batch_size" not in self._freeze_dims and any( + dimension.dim_param == "batch_size" + for value in model.graph.input + if value.name not in initializer_names and value.type.HasField("tensor_type") + for dimension in value.type.tensor_type.shape.dim + ): + self._freeze_dims["batch_size"] = 1 + self.console.print("[dim]Freeze input dimension: batch_size=1[/dim]") + + def export( + self, + model: str | Path, + output_path: str | Path, + ) -> None: + """Export ONNX to CGC; retained as a compatibility alias.""" + self.export_onnx(model=model, output_path=output_path) + + def output_artifacts(self, output_path: str | Path) -> tuple[Path, ...]: + """Return primary and potential sidecar artifacts for output guarding.""" + output_path = Path(output_path) + artifacts = [output_path] + if self.options.external_weights: + artifacts.append(output_path.with_name(f"{output_path.name}.data")) + artifacts.append(cgc_metadata_path(output_path)) + return tuple(artifacts) + + def _write_io_metadata( + self, + result: CGCExportResult, + output_path: Path, + ) -> None: + """Persist source ONNX names for ordinal-only Runtime artifacts.""" + cgc_metadata_path(output_path).write_text( + json.dumps( + { + "format": "cgc-input-ir", + "inputs": [ + {"name": name, "index": index} + for index, name in enumerate(result.input_names) + ], + "outputs": [ + {"name": name, "index": index} + for index, name in enumerate(result.output_names) + ], + }, + indent=2, + ), + encoding="utf-8", + ) + + @staticmethod + def _publish_onnx_reports( + intermediate_path: Path, + output_path: Path, + ) -> None: + """Move requested HTP reports out of the temporary ONNX directory.""" + intermediate_base = intermediate_path.with_suffix("") + output_base = output_path.with_suffix("") + for suffix in ("_htp_metadata.json", "_htp_export_report.md"): + source = intermediate_base.with_name(f"{intermediate_base.name}{suffix}") + destination = output_base.with_name(f"{output_base.name}{suffix}") + source.replace(destination) + + def _export_mlir( + self, + source_path: Path, + output_path: Path, + result: CGCExportResult, + ) -> None: + """Generate standalone CGC MLIR with optional Foundry weight data.""" + final_weights = ( + output_path.with_name(f"{output_path.name}.data") + if self.options.external_weights + else None + ) + final_metadata = cgc_metadata_path(output_path) + with tempfile.TemporaryDirectory( + prefix=f".{output_path.name}.", + dir=output_path.parent, + ) as temporary: + staging_dir = Path(temporary) + staged_output = staging_dir / output_path.name + staged_weights = ( + staging_dir / final_weights.name if final_weights is not None else None + ) + staged_metadata = cgc_metadata_path(staged_output) + + ir = self._convert_to_cgir( + source_path, + output_data_file=staged_weights, + ) + staged_output.write_text(ir, encoding="utf-8", newline="\n") + self._write_io_metadata(result, staged_output) + + artifacts = [(staged_metadata, final_metadata), (staged_output, output_path)] + if staged_weights is not None and final_weights is not None: + artifacts.insert(0, (staged_weights, final_weights)) + for _, destination in artifacts: + if destination.exists() and not destination.is_file(): + raise ValueError(f"Output artifact exists but is not a file: {destination}") + + backup_dir = staging_dir / "backup" + backup_dir.mkdir() + backups: list[tuple[Path, Path]] = [] + published: list[Path] = [] + try: + for index, (_, destination) in enumerate(artifacts): + if destination.exists() or destination.is_symlink(): + backup = backup_dir / str(index) + destination.replace(backup) + backups.append((backup, destination)) + for staged, destination in artifacts: + if staged.exists(): + staged.replace(destination) + published.append(destination) + except BaseException: + for destination in reversed(published): + destination.unlink(missing_ok=True) + for backup, destination in reversed(backups): + backup.replace(destination) + raise + + def _convert_to_cgir( + self, + source_path: Path, + *, + output_data_file: Path | None = None, + ) -> str: + """Compile ONNX to textual CGC Input IR with FoundryToolbox.""" + try: + with FoundryCompiler() as compiler: + # model_directory lets Foundry resolve source ONNX external-data + # locations while lazy external mode avoids loading those weights. + serialized = compiler.compile_onnx( + source_path.read_bytes(), + model_directory=source_path.parent, + update_opset=self.options.update_opset, + topo_sort_nodes=self.options.topo_sort_nodes, + include_initializers=not self.options.external_weights, + enable_lazy_external_data=self.options.external_weights, + output_data_file=output_data_file, + freeze_dims=self._freeze_dims, + ) + except FoundryCompileError as e: + raise RuntimeError(self._format_foundry_error(e)) from e + + try: + ir = serialized.decode("utf-8") + except UnicodeDecodeError as e: + raise RuntimeError( + "FoundryToolbox returned non-UTF-8 CGC textual MLIR" + ) from e + if not ir.strip(): + raise RuntimeError("FoundryToolbox returned empty CGC textual MLIR") + if "cgc." not in ir.lower() and "#cgc" not in ir.lower(): + raise RuntimeError("FoundryToolbox output does not contain the CGC dialect") + return ir + + def _format_foundry_error(self, error: FoundryCompileError) -> str: + """Format native Foundry diagnostics for an export user.""" + details = [f"Foundry failed to convert ONNX to CGC IR [{error.result_name}]."] + if error.unsupported_op: + details.append(f"ONNX operator '{error.unsupported_op}' is not supported.") + elif error.missing_external_data: + details.append( + f"ONNX external weights file was not found: " + f"'{error.missing_external_data}'." + ) + if error.native_message: + details.append(error.native_message) + if error.result_name == "SHAPE_INFERENCE" and not self._freeze_dims: + details.append( + "If symbolic dimensions caused this failure, retry with " + "--options freeze-dims=batch=1,seq=128 " + "using the model's dimension names." + ) + return " ".join(details) + +def export_cgc( + model: str | Path, + output_path: str | Path, + options: CGCOptions, +) -> None: + """Export an ONNX model using :class:`CGCExporter`.""" + CGCExporter(options).export_onnx( + model=model, + output_path=output_path, + ) + + +def _parse_freeze_dims(value: str) -> dict[str, int]: + """Parse comma-separated symbolic dimension overrides.""" + if not value.strip(): + return {} + overrides: dict[str, int] = {} + for assignment in value.split(","): + name, separator, raw_size = assignment.partition("=") + name = name.strip() + raw_size = raw_size.strip() + if not separator or not name or not raw_size: + raise ValueError( + "freeze-dims expects comma-separated NAME=SIZE values, " + f"got {assignment!r}." + ) + try: + size = int(raw_size) + except ValueError as e: + raise ValueError( + f"freeze-dims value for {name!r} must be an integer, got {raw_size!r}." + ) from e + if size <= 0: + raise ValueError( + f"freeze-dims value for {name!r} must be greater than zero." + ) + if name in overrides: + raise ValueError(f"freeze-dims contains duplicate dimension {name!r}.") + overrides[name] = size + return overrides + + +__all__ = ["CGCExportResult", "CGCExporter", "CGCOptions", "export_cgc"] diff --git a/src/winml/modelkit/export/cgc/foundry.py b/src/winml/modelkit/export/cgc/foundry.py new file mode 100644 index 000000000..3a4fb4b28 --- /dev/null +++ b/src/winml/modelkit/export/cgc/foundry.py @@ -0,0 +1,336 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Lazy ctypes binding for the FoundryToolbox compiler C API.""" + +from __future__ import annotations + +import ctypes +import os +from importlib import metadata +from pathlib import Path +from typing import TYPE_CHECKING + + +if TYPE_CHECKING: + from collections.abc import Mapping + + +FDY_VERSION = 0x00000004 +FDY_SOURCE_FORMAT_ONNX_PROTOBUF = 0 +FDY_COMPILER_TARGET_DXCGC = 1 +FDY_SERIALIZATION_FORMAT_TEXT = 1 +FDY_COMPILER_RESULT_SUCCESS = 0 +FDY_PASS_OVERRIDE_DYNAMIC_DIMS_BY_DIM_NAME = 15 +FDY_PASS_STAGE_BEFORE_LOWERING = 0 + +_RESULT_NAMES = { + 1: "INVALID_ARGUMENT", + 2: "PARSE", + 3: "SHAPE_INFERENCE", + 4: "IMPORT", + 5: "VERIFICATION", + 6: "LOWERING", + 7: "SERIALIZATION", + 8: "BUFFER_TOO_SMALL", + 9: "UNSUPPORTED_OP", + 10: "MALFORMED", + 11: "MISSING_EXTERNAL_DATA", +} + + +class FoundryToolboxUnavailableError(RuntimeError): + """Raised when a compatible FoundryToolbox DLL cannot be located or loaded.""" + + +class FoundryCompileError(RuntimeError): + """Native Foundry failure with copied diagnostic fields.""" + + def __init__( + self, + result_code: int, + *, + native_message: str = "", + unsupported_op: str | None = None, + missing_external_data: str | None = None, + ) -> None: + self.result_code = result_code + self.result_name = _RESULT_NAMES.get(result_code, "UNKNOWN") + self.native_message = native_message + self.unsupported_op = unsupported_op + self.missing_external_data = missing_external_data + suffix = f": {native_message}" if native_message else "" + super().__init__(f"Foundry compiler failed [{self.result_name}]{suffix}") + + +class _FdyStringView(ctypes.Structure): + _fields_ = [("data", ctypes.c_char_p), ("size", ctypes.c_size_t)] + + +class _FdySpan(ctypes.Structure): + _fields_ = [("data", ctypes.c_void_p), ("size", ctypes.c_size_t)] + + +class _FdyMutableSpan(ctypes.Structure): + _fields_ = [("data", ctypes.c_void_p), ("size", ctypes.c_size_t)] + + +class _FdyPassDescriptor(ctypes.Structure): + _fields_ = [("kind", ctypes.c_uint32), ("stage", ctypes.c_uint32)] + + +class _FdyOverrideDynamicDimsByDimNamePassDescriptor(ctypes.Structure): + _fields_ = [ + ("descriptor", _FdyPassDescriptor), + ("names", ctypes.POINTER(ctypes.c_char_p)), + ("values", ctypes.POINTER(ctypes.c_int64)), + ("count", ctypes.c_size_t), + ] + + +class _FdyCompilerOptions(ctypes.Structure): + _fields_ = [ + ("version", ctypes.c_uint32), + ("sourceFormat", ctypes.c_uint32), + ("target", ctypes.c_uint32), + ("updateOpset", ctypes.c_bool), + ("topoSortNodes", ctypes.c_bool), + ("includeInitializers", ctypes.c_bool), + ("outputDataFile", _FdyStringView), + ("passes", ctypes.POINTER(ctypes.POINTER(_FdyPassDescriptor))), + ("passCount", ctypes.c_uint32), + ("modelDirectory", _FdyStringView), + ("enableLazyExternalData", ctypes.c_bool), + ("safetensorsFiles", ctypes.POINTER(_FdyStringView)), + ("safetensorsFileCount", ctypes.c_uint32), + ] + + +def find_foundry_toolbox() -> Path: + """Resolve FoundryToolbox from the installed windowsml wheel.""" + try: + distribution = metadata.distribution("windowsml") + except metadata.PackageNotFoundError as e: + raise FoundryToolboxUnavailableError( + "CGC export requires a windowsml wheel containing FoundryToolbox.dll." + ) from e + + candidate = Path( + str(distribution.locate_file("windowsml/lib/FoundryToolbox.dll")) + ) + if not candidate.is_file(): + raise FoundryToolboxUnavailableError( + "The installed windowsml wheel " + f"({distribution.version}) does not contain FoundryToolbox.dll." + ) + return candidate.resolve() + + +def _string_view(value: bytes) -> _FdyStringView: + return _FdyStringView(value or None, len(value)) + + +class FoundryCompiler: + """Thin owner for one Foundry compiler context.""" + + def __init__(self) -> None: + path = find_foundry_toolbox() + self._dll_directory = None + try: + if hasattr(os, "add_dll_directory"): + self._dll_directory = os.add_dll_directory(str(path.parent)) + self._dll = ctypes.CDLL(str(path)) + except OSError as e: + if self._dll_directory is not None: + self._dll_directory.close() + raise FoundryToolboxUnavailableError( + f"Unable to load FoundryToolbox.dll from '{path}': {e}" + ) from e + + self._compiler = ctypes.c_void_p() + try: + self._configure_signatures() + result = self._dll.FdyCompilerCreate(ctypes.byref(self._compiler)) + except Exception: + self.close() + raise + if result != FDY_COMPILER_RESULT_SUCCESS: + self.close() + raise FoundryCompileError( + result, + native_message="Foundry compiler context creation failed.", + ) + + def _configure_signatures(self) -> None: + dll = self._dll + dll.FdyCompilerCreate.argtypes = [ctypes.POINTER(ctypes.c_void_p)] + dll.FdyCompilerCreate.restype = ctypes.c_uint32 + dll.FdyCompilerDestroy.argtypes = [ctypes.c_void_p] + dll.FdyCompilerDestroy.restype = None + dll.FdyCompilerCompile.argtypes = [ + ctypes.c_void_p, + _FdySpan, + ctypes.POINTER(_FdyCompilerOptions), + ctypes.POINTER(ctypes.c_void_p), + ] + dll.FdyCompilerCompile.restype = ctypes.c_uint32 + dll.FdyModuleSerialize.argtypes = [ + ctypes.c_void_p, + ctypes.c_uint32, + _FdyMutableSpan, + ctypes.POINTER(ctypes.c_size_t), + ] + dll.FdyModuleSerialize.restype = ctypes.c_uint32 + dll.FdyModuleDestroy.argtypes = [ctypes.c_void_p] + dll.FdyModuleDestroy.restype = None + for name in ( + "FdyCompilerGetLastError", + "FdyCompilerGetLastUnsupportedOpName", + "FdyCompilerGetLastMissingExternalDataFile", + ): + function = getattr(dll, name) + function.argtypes = [ctypes.c_void_p] + function.restype = ctypes.c_char_p + + def close(self) -> None: + """Release the native compiler context.""" + if getattr(self, "_compiler", None): + self._dll.FdyCompilerDestroy(self._compiler) + self._compiler = ctypes.c_void_p() + if self._dll_directory is not None: + self._dll_directory.close() + self._dll_directory = None + + def __enter__(self) -> FoundryCompiler: + return self + + def __exit__(self, *_args: object) -> None: + self.close() + + def _last_error(self, result: int) -> FoundryCompileError: + # Foundry owns these strings only until the next call on this compiler. + # Decode each value immediately so the exception owns durable copies. + raw_message = self._dll.FdyCompilerGetLastError(self._compiler) + message = raw_message.decode("utf-8", errors="replace") if raw_message else "" + unsupported_op = None + missing_external_data = None + if result == 9: + raw_detail = self._dll.FdyCompilerGetLastUnsupportedOpName(self._compiler) + if raw_detail: + unsupported_op = raw_detail.decode("utf-8", errors="replace") + elif result == 11: + raw_detail = self._dll.FdyCompilerGetLastMissingExternalDataFile(self._compiler) + if raw_detail: + missing_external_data = raw_detail.decode("utf-8", errors="replace") + return FoundryCompileError( + result, + native_message=message, + unsupported_op=unsupported_op, + missing_external_data=missing_external_data, + ) + + def compile_onnx( + self, + source: bytes, + *, + model_directory: Path, + update_opset: bool, + topo_sort_nodes: bool, + include_initializers: bool, + enable_lazy_external_data: bool, + output_data_file: Path | None = None, + freeze_dims: Mapping[str, int] | None = None, + ) -> bytes: + """Compile ONNX protobuf bytes to textual DxCGC MLIR.""" + source_buffer = ctypes.create_string_buffer(source) + model_directory_bytes = str(model_directory.resolve()).encode("utf-8") + output_data_file_bytes = ( + str(output_data_file).encode("utf-8") + if output_data_file is not None + else b"" + ) + dim_names = tuple((freeze_dims or {}).keys()) + dim_name_bytes = tuple(name.encode("utf-8") for name in dim_names) + dim_name_array = ( + (ctypes.c_char_p * len(dim_names))(*dim_name_bytes) + if dim_names + else None + ) + dim_value_array = ( + (ctypes.c_int64 * len(dim_names))( + *((freeze_dims or {})[name] for name in dim_names) + ) + if dim_names + else None + ) + dim_descriptor = ( + _FdyOverrideDynamicDimsByDimNamePassDescriptor( + descriptor=_FdyPassDescriptor( + kind=FDY_PASS_OVERRIDE_DYNAMIC_DIMS_BY_DIM_NAME, + stage=FDY_PASS_STAGE_BEFORE_LOWERING, + ), + names=dim_name_array, + values=dim_value_array, + count=len(dim_names), + ) + if dim_names + else None + ) + passes = ( + (ctypes.POINTER(_FdyPassDescriptor) * 1)( + ctypes.cast( + ctypes.pointer(dim_descriptor), + ctypes.POINTER(_FdyPassDescriptor), + ) + ) + if dim_descriptor is not None + else None + ) + compiler_options = _FdyCompilerOptions( + version=FDY_VERSION, + sourceFormat=FDY_SOURCE_FORMAT_ONNX_PROTOBUF, + target=FDY_COMPILER_TARGET_DXCGC, + updateOpset=update_opset, + topoSortNodes=topo_sort_nodes, + includeInitializers=include_initializers, + outputDataFile=_string_view(output_data_file_bytes), + passes=passes, + passCount=len(passes) if passes is not None else 0, + modelDirectory=_string_view(model_directory_bytes), + enableLazyExternalData=enable_lazy_external_data, + safetensorsFiles=None, + safetensorsFileCount=0, + ) + module = ctypes.c_void_p() + result = self._dll.FdyCompilerCompile( + self._compiler, + _FdySpan(ctypes.cast(source_buffer, ctypes.c_void_p), len(source)), + ctypes.byref(compiler_options), + ctypes.byref(module), + ) + if result != FDY_COMPILER_RESULT_SUCCESS: + raise self._last_error(result) + + try: + size = ctypes.c_size_t() + result = self._dll.FdyModuleSerialize( + module, + FDY_SERIALIZATION_FORMAT_TEXT, + _FdyMutableSpan(None, 0), + ctypes.byref(size), + ) + if result != FDY_COMPILER_RESULT_SUCCESS: + raise self._last_error(result) + output = ctypes.create_string_buffer(size.value) + result = self._dll.FdyModuleSerialize( + module, + FDY_SERIALIZATION_FORMAT_TEXT, + _FdyMutableSpan(ctypes.cast(output, ctypes.c_void_p), len(output)), + ctypes.byref(size), + ) + if result != FDY_COMPILER_RESULT_SUCCESS: + raise self._last_error(result) + return output.raw[: size.value] + finally: + self._dll.FdyModuleDestroy(module) diff --git a/src/winml/modelkit/export/config.py b/src/winml/modelkit/export/config.py index 67ce7ff18..f68fad4a2 100644 --- a/src/winml/modelkit/export/config.py +++ b/src/winml/modelkit/export/config.py @@ -18,6 +18,7 @@ # InputTensorSpec and OutputTensorSpec live in modelkit.onnx.io (canonical home). from ..onnx import InputTensorSpec, OutputTensorSpec +from ..utils.constants import EXPORT_TARGETS, ExportTarget from .policy import ExportCompatibilityConfig @@ -195,6 +196,9 @@ class WinMLExportConfig: input_names_: InitVar[list[str] | None] = None output_names_: InitVar[list[str] | None] = None + target: ExportTarget = "onnx" + options: dict[str, Any] = field(default_factory=dict) + def __post_init__( self, input_shape_: tuple[int, ...] | None, @@ -202,6 +206,12 @@ def __post_init__( output_names_: list[str] | None, ) -> None: """Validate configuration after initialization.""" + if self.target not in EXPORT_TARGETS: + raise ValueError( + f"Invalid export target {self.target!r}. Must be one of {EXPORT_TARGETS}" + ) + if not isinstance(self.options, dict): + raise TypeError("Export options must be a JSON object") # Handle legacy parameters - convert to input_tensors/output_tensors if needed if input_shape_ is not None and self.input_tensors is None: # Convert legacy input_shape to input_tensors @@ -375,6 +385,11 @@ def to_dict(self) -> dict[str, Any]: if self.compatibility: result["compatibility"] = self.compatibility.to_dict() + if self.target != "onnx": + result["target"] = self.target + if self.options: + result["options"] = dict(self.options) + return result @classmethod @@ -387,6 +402,9 @@ def from_dict(cls, data: dict[str, Any]) -> WinMLExportConfig: Returns: WinMLExportConfig instance. """ + if not isinstance(data, dict): + raise TypeError("Export configuration must be a JSON object") + # Parse input_tensors if present input_tensors = None raw_inputs = data.get("input_tensors") @@ -406,6 +424,8 @@ def from_dict(cls, data: dict[str, Any]) -> WinMLExportConfig: ] return cls( + target=data.get("target", "onnx"), + options=data.get("options", {}), opset_version=data.get("opset_version", 17), batch_size=data.get("batch_size", 1), input_tensors=input_tensors, diff --git a/src/winml/modelkit/models/auto.py b/src/winml/modelkit/models/auto.py index 907d91c86..2991290db 100644 --- a/src/winml/modelkit/models/auto.py +++ b/src/winml/modelkit/models/auto.py @@ -33,6 +33,7 @@ from ..config import WinMLBuildConfig from ..loader.task import get_task_abbrev from ..session import short_ep_name +from ..utils.constants import resolve_runtime_api_backend # Import task mapping from winml/ subpackage from .winml import get_supported_tasks, get_winml_class @@ -45,6 +46,7 @@ from ..build import BuildResult from ..session import WinMLEPDevice + from ..utils.constants import RuntimeBackend, RuntimeName from .winml.base import WinMLPreTrainedModel from .winml.composite_model import WinMLCompositeModel @@ -53,11 +55,14 @@ def _get_cache_build_controls( *, + skip_build: bool = False, skip_optimize: bool = False, hack_max_optim_iterations: int | None = None, ) -> dict[str, Any]: """Return only the non-default artifact-changing build controls.""" build_controls: dict[str, Any] = {} + if skip_build: + build_controls["skip_build"] = True if skip_optimize: build_controls["skip_optimize"] = True if hack_max_optim_iterations is not None and hack_max_optim_iterations != 3: @@ -73,6 +78,15 @@ def _resolved_ep_short_name(ep_device: WinMLEPDevice) -> str: return short_ep_name(ep_device.device.ep_name) +def _uses_cgc_online( + runtime: RuntimeName, backend: RuntimeBackend | None, ep_device: WinMLEPDevice +) -> bool: + """Return whether the execution layer handles CGC conversion from ONNX.""" + return (runtime == "winml-runtime" and backend == "cgc") or ( + runtime == "winml-ort" and _resolved_ep_short_name(ep_device) == "winmlcg" + ) + + @dataclass(frozen=True) class _PretrainedArtifact: result: "BuildResult" @@ -153,6 +167,8 @@ def from_onnx( compile_provider_options: dict[str, str] | None = None, session_options: Callable[[], Any] | None = None, hf_config: PretrainedConfig | None = None, + runtime: RuntimeName = "winml-ort", + backend: RuntimeBackend | None = None, **kwargs: Any, ) -> WinMLPreTrainedModel | WinMLCompositeModel: """Build from a pre-exported ONNX file. @@ -179,12 +195,15 @@ def from_onnx( Returns: WinMLPreTrainedModel inference wrapper. """ + backend = resolve_runtime_api_backend(runtime, onnx_path, backend) + # Ergonomic path: resolve ep_device from device/ep shortcuts. if ep_device is None: from ..session import EPDeviceTarget, WinMLEPRegistry, resolve_device target = resolve_device( - EPDeviceTarget(ep=ep or "auto", device=(device or "auto").lower()) + EPDeviceTarget(ep=ep or "auto", device=(device or "auto").lower()), + backend=backend, ) ep_device = WinMLEPRegistry.instance().auto_device(target) @@ -205,6 +224,8 @@ def from_onnx( provider_options=provider_options, compile_provider_options=compile_provider_options, session_options=session_options, + runtime=runtime, + backend=backend, **kwargs, ) @@ -219,6 +240,7 @@ def from_onnx( # If user provides config, treat it as an override (merged on top). from ..config import generate_onnx_build_config + skip_build = skip_build or _uses_cgc_online(runtime, backend, ep_device) config = generate_onnx_build_config( onnx_path, task=task, @@ -227,8 +249,10 @@ def from_onnx( ep=_resolved_ep_short_name(ep_device), override=config, no_compile=no_compile, + backend=backend, ) - if compile_provider_options: + + if compile_provider_options and not skip_build: if config.compile is None: raise ValueError("compile_provider_options requires compilation to be enabled.") config.compile.ep_config.provider_options = { @@ -254,6 +278,8 @@ def from_onnx( ep_device=ep_device, provider_options=provider_options, session_options=session_options, + runtime=runtime, + backend=backend, ) # Resolve output directory @@ -308,6 +334,38 @@ def from_onnx( ep_device=ep_device, provider_options=provider_options, session_options=session_options, + runtime=runtime, + backend=backend, + ) + + @classmethod + def from_mlir( + cls, + mlir_path: str | Path, + *, + ep_device: WinMLEPDevice, + task: str | None = None, + runtime: RuntimeName = "winml-runtime", + backend: RuntimeBackend = "cgc", + ) -> WinMLPreTrainedModel: + """Load a pre-built CGC MLIR artifact.""" + if runtime != "winml-runtime": + raise ValueError("MLIR inputs require runtime='winml-runtime'.") + if task == "text-generation": + raise ValueError("from_mlir does not support task='text-generation'.") + resolved_backend = resolve_runtime_api_backend(runtime, mlir_path, backend) + + mlir_path = Path(mlir_path) + if not mlir_path.is_file(): + raise FileNotFoundError(f"CGC MLIR model not found: {mlir_path}") + + winml_class = get_winml_class(None, task) + return winml_class( + onnx_path=mlir_path, + config=None, + ep_device=ep_device, + runtime=runtime, + backend=resolved_backend, ) @classmethod @@ -324,6 +382,7 @@ def from_pretrained( cache_dir: str | Path | None = None, use_cache: bool = True, force_rebuild: bool = False, + skip_build: bool = False, trust_remote_code: bool = False, shape_config: dict | None = None, model_type: str | None = None, @@ -333,6 +392,8 @@ def from_pretrained( no_compile: bool = False, skip_optimize: bool = False, hack_max_optim_iterations: int = 3, + runtime: RuntimeName = "winml-ort", + backend: RuntimeBackend | None = None, **kwargs: Any, ) -> WinMLPreTrainedModel | WinMLCompositeModel: """Load appropriate WinML model based on task detection. @@ -360,6 +421,7 @@ def from_pretrained( use_cache: If True (default), use persistent cache directory. If False, build in a temp directory and always rebuild. force_rebuild: If True, rebuild even if cached model exists. + skip_build: Use the original ONNX or HF export without further build stages. trust_remote_code: Whether to trust remote code in HF models shape_config: Shape overrides passed to generate_build_config(). Valid keys -- text: sequence_length; vision: height, width; @@ -378,6 +440,7 @@ def from_pretrained( model_input = resolve_model_input(str(model_id_or_path)) model_id = model_input.local_path or model_input.raw + backend = resolve_runtime_api_backend(runtime, model_id, backend) logger.info("Loading WinML model from: %s", model_id) request_device = (device or "auto").lower() request_ep = ep @@ -388,7 +451,8 @@ def from_pretrained( from ..session import EPDeviceTarget, WinMLEPRegistry, resolve_device target = resolve_device( - EPDeviceTarget(ep=ep or "auto", device=(device or "auto").lower()) + EPDeviceTarget(ep=ep or "auto", device=(device or "auto").lower()), + backend=backend, ) ep_device = WinMLEPRegistry.instance().auto_device(target) @@ -408,9 +472,12 @@ def from_pretrained( cache_dir=cache_dir, use_cache=use_cache, force_rebuild=force_rebuild, + skip_build=skip_build, no_compile=no_compile, provider_options=provider_options, session_options=session_options, + runtime=runtime, + backend=backend, allow_unsupported_nodes=allow_unsupported_nodes, skip_optimize=skip_optimize, hack_max_optim_iterations=hack_max_optim_iterations, @@ -462,8 +529,11 @@ def from_pretrained( device=request_device, ep=request_ep, ep_device=ep_device, + runtime=runtime, + backend=backend, use_cache=use_cache, force_rebuild=force_rebuild, + skip_build=skip_build, trust_remote_code=trust_remote_code, shape_config=shape_config, precision=precision, @@ -492,6 +562,7 @@ def from_pretrained( cache_dir=cache_dir, use_cache=use_cache, force_rebuild=force_rebuild, + skip_build=skip_build, trust_remote_code=trust_remote_code, shape_config=shape_config, model_type=model_type, @@ -499,6 +570,8 @@ def from_pretrained( no_compile=no_compile, skip_optimize=skip_optimize, hack_max_optim_iterations=hack_max_optim_iterations, + runtime=runtime, + backend=backend, **kwargs, ) onnx_path = artifact.result.final_onnx_path @@ -515,6 +588,8 @@ def from_pretrained( ep_device=ep_device, provider_options=provider_options, session_options=session_options, + runtime=runtime, + backend=backend, ) model._build_config = artifact.build_config return model @@ -533,6 +608,7 @@ def _build_pretrained_artifact( cache_dir: str | Path | None = None, use_cache: bool = True, force_rebuild: bool = False, + skip_build: bool = False, trust_remote_code: bool = False, shape_config: dict | None = None, model_type: str | None = None, @@ -540,22 +616,29 @@ def _build_pretrained_artifact( no_compile: bool = False, skip_optimize: bool = False, hack_max_optim_iterations: int = 3, + runtime: RuntimeName = "winml-ort", + backend: RuntimeBackend | None = None, **_kwargs: Any, ) -> _PretrainedArtifact: from ..utils.model_input import resolve_model_input model_input = resolve_model_input(str(model_id_or_path)) model_id = model_input.local_path or model_input.raw + backend = resolve_runtime_api_backend(runtime, model_id, backend) request_device = (device or "auto").lower() request_ep = ep if ep_device is None: from ..session import EPDeviceTarget, WinMLEPRegistry, resolve_device - target = resolve_device(EPDeviceTarget(ep=request_ep or "auto", device=request_device)) + target = resolve_device( + EPDeviceTarget(ep=request_ep or "auto", device=request_device), + backend=backend, + ) ep_device = WinMLEPRegistry.instance().auto_device(target) runtime_device = ep_device.device.device_type.lower() runtime_ep = _resolved_ep_short_name(ep_device) + skip_build = skip_build or _uses_cgc_online(runtime, backend, ep_device) from ..config import generate_hf_build_config @@ -572,6 +655,7 @@ def _build_pretrained_artifact( trust_remote_code=trust_remote_code, policy_overrides_config=True, no_compile=no_compile, + backend=backend, ) resolved_task = cast("str", build_config.loader.task) @@ -605,6 +689,7 @@ def _build_pretrained_artifact( get_task_abbrev(resolved_task), build_config.generate_cache_key(), _get_cache_build_controls( + skip_build=skip_build, skip_optimize=skip_optimize, hack_max_optim_iterations=hack_max_optim_iterations, ), @@ -623,6 +708,7 @@ def _build_pretrained_artifact( output_dir=output_dir, model_id=model_id, rebuild=force_rebuild, + skip_build=skip_build, trust_remote_code=trust_remote_code, cache_key=cache_key, ep=resolved_ep, diff --git a/src/winml/modelkit/models/winml/base.py b/src/winml/modelkit/models/winml/base.py index d3e1d341a..ca5a53299 100644 --- a/src/winml/modelkit/models/winml/base.py +++ b/src/winml/modelkit/models/winml/base.py @@ -31,6 +31,9 @@ import contextlib from collections.abc import Callable + from ...utils.constants import RuntimeBackend, RuntimeName + +from ...session.runtime_session import WinMLRuntimeSession from ...session.session import WinMLSession @@ -41,6 +44,11 @@ logger = logging.getLogger(__name__) +SESSION_CLASSES: dict[str, type[WinMLSession] | type[WinMLRuntimeSession]] = { + "winml-ort": WinMLSession, + "winml-runtime": WinMLRuntimeSession, +} + class PreTrainedModel: """Name shim so HF ``infer_framework()`` recognizes WinML models as "pt". @@ -65,10 +73,12 @@ class WinMLPreTrainedModel(PreTrainedModel, ABC): def __init__( self, onnx_path: str | Path, - ep_device: WinMLEPDevice, + ep_device: WinMLEPDevice | None, config: PretrainedConfig | None = None, provider_options: dict[str, str] | None = None, session_options: Callable[[], Any] | None = None, + runtime: RuntimeName = "winml-ort", + backend: RuntimeBackend | None = None, ) -> None: """Initialize inference model. @@ -85,12 +95,13 @@ def __init__( # Set by WinMLAutoModel.from_pretrained() after construction self._build_config: Any = None - # Create WinMLSession (delegates ORT operations) - self._session = WinMLSession( - onnx_path=self._onnx_path, + runtime_kwargs: dict[str, Any] = {"backend": backend} if runtime == "winml-runtime" else {} + self._session = SESSION_CLASSES[runtime]( + self._onnx_path, ep_device=ep_device, provider_options=provider_options, session_options=session_options, + **runtime_kwargs, ) @property diff --git a/src/winml/modelkit/models/winml/composite_model.py b/src/winml/modelkit/models/winml/composite_model.py index 1455ba737..38b0ec2cf 100644 --- a/src/winml/modelkit/models/winml/composite_model.py +++ b/src/winml/modelkit/models/winml/composite_model.py @@ -55,6 +55,7 @@ from transformers import PretrainedConfig from ...session import WinMLEPDevice + from ...utils.constants import RuntimeName logger = logging.getLogger(__name__) @@ -135,6 +136,7 @@ def from_pretrained( force_rebuild: bool = False, sub_model_kwargs: dict[str, dict[str, Any]] | None = None, trust_remote_code: bool = False, + runtime: RuntimeName = "winml-ort", **kwargs: Any, ) -> WinMLCompositeModel: """Build all sub-components and return ready-to-use model. @@ -202,6 +204,7 @@ def from_pretrained( force_rebuild=force_rebuild, sub_model_kwargs=sub_model_kwargs, trust_remote_code=trust_remote_code, + runtime=runtime, **kwargs, ) from ..auto import WinMLAutoModel @@ -228,6 +231,7 @@ def from_pretrained( use_cache=use_cache, force_rebuild=force_rebuild, trust_remote_code=trust_remote_code, + runtime=runtime, **merged, ) @@ -243,6 +247,7 @@ def from_onnx( task: str | None = None, hf_config: PretrainedConfig | None = None, sub_model_kwargs: dict[str, dict[str, Any]] | None = None, + runtime: RuntimeName = "winml-ort", **kwargs: Any, ) -> WinMLCompositeModel: """Load composite model from pre-built ONNX files. @@ -296,7 +301,11 @@ def from_onnx( f"Unknown component {name!r}. Valid names for {resolved_cls.__name__}: {valid}" ) merged = {**kwargs, "task": component_task, **per_component.get(name, {})} - sub_models[name] = WinMLAutoModel.from_onnx(Path(path), **merged) + sub_models[name] = WinMLAutoModel.from_onnx( + Path(path), + runtime=runtime, + **merged, + ) if hf_config is None: raise ValueError("Composite model construction requires an HF config (hf_config).") diff --git a/src/winml/modelkit/optim/capabilities/graph.py b/src/winml/modelkit/optim/capabilities/graph.py index cab11d679..b1ed26e2b 100644 --- a/src/winml/modelkit/optim/capabilities/graph.py +++ b/src/winml/modelkit/optim/capabilities/graph.py @@ -14,6 +14,15 @@ from ..registry import BoolCapability, CapabilityCategory +ORT_GRAPH_OPTIMIZATION = BoolCapability( + name="ort-graph-optimization", + ort_name=None, + description="Run ORT graph optimization, including basic optimizations and graph fusions", + category=CapabilityCategory.GRAPH, + default=True, +) + + # Concat-slice elimination - remove concat followed by slice CONCAT_SLICE_ELIMINATION = BoolCapability( name="concat-slice-elimination", diff --git a/src/winml/modelkit/optim/config.py b/src/winml/modelkit/optim/config.py index e85e6ae19..541835c95 100644 --- a/src/winml/modelkit/optim/config.py +++ b/src/winml/modelkit/optim/config.py @@ -21,6 +21,16 @@ class WinMLOptimizationConfig(dict): def __init__(self, **kwargs: bool) -> None: super().__init__(kwargs) + @classmethod + def for_cgc(cls) -> WinMLOptimizationConfig: + """Enable the registered CGIR compatibility rules without ORT graph optimization.""" + from .pipes import CGIRRewritePipe + + return cls( + ort_graph_optimization=False, + **CGIRRewritePipe.get_compatibility_options(), + ) + def to_dict(self) -> dict: """Convert to dictionary (sorted keys for deterministic serialization).""" return dict(sorted(self.items())) diff --git a/src/winml/modelkit/optim/pipes/__init__.py b/src/winml/modelkit/optim/pipes/__init__.py index 9315b71fa..54a76f6e0 100644 --- a/src/winml/modelkit/optim/pipes/__init__.py +++ b/src/winml/modelkit/optim/pipes/__init__.py @@ -17,6 +17,7 @@ AlgebraicRewritePipeConfig, ) from .base import BasePipe, OptimizationError, PipeConfig, caps_dict +from .cgir_rewrite import CGIRRewritePipe, CGIRRewritePipeConfig from .fusion import ORTFusionPipe, ORTFusionPipeConfig from .graph import GRAPH_CAPABILITIES, ORTGraphPipe, ORTGraphPipeConfig from .rewrite import RewritePipe, RewritePipeConfig @@ -24,9 +25,10 @@ # Optimization pipes to run in sequence +# - CGIRRewritePipe: Explicit target-specific rewrites before an EP can compile the graph. # - ORTGraphPipe: ORT graph-level optimizations (C++ optimizer), including constant folding. -# Runs first so downstream pipes see a constant-folded graph (e.g. Reshape shape inputs -# become literal constants, enabling skeleton-based pattern matching). +# Runs before general downstream pipes so they see a constant-folded graph (e.g. Reshape +# shape inputs become literal constants, enabling skeleton-based pattern matching). # - AlgebraicRewritePipe: Exact topology-based algebraic rewrites (after ORT folding). # - RewritePipe: Pattern-based subgraph rewriting (runs after ORT constant folding so that # shape constants are visible, but before ORTFusionPipe so normalised patterns are @@ -34,6 +36,7 @@ # - ORTFusionPipe: ORT transformer fusions (Python optimizer) # - SurgeryPipe: Post-optimization model surgery (runs last to clamp constants after folding) PIPES: list[type[BasePipe]] = [ + CGIRRewritePipe, ORTGraphPipe, AlgebraicRewritePipe, RewritePipe, @@ -62,6 +65,8 @@ def get_all_capabilities() -> dict[str, Any]: "AlgebraicRewritePipe", "AlgebraicRewritePipeConfig", "BasePipe", + "CGIRRewritePipe", + "CGIRRewritePipeConfig", "ORTFusionPipe", "ORTFusionPipeConfig", "ORTGraphPipe", diff --git a/src/winml/modelkit/optim/pipes/cgir_rewrite.py b/src/winml/modelkit/optim/pipes/cgir_rewrite.py new file mode 100644 index 000000000..e28ab6fd4 --- /dev/null +++ b/src/winml/modelkit/optim/pipes/cgir_rewrite.py @@ -0,0 +1,183 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""CGIR-specific model rewrites applied before target EP compilation.""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any, ClassVar + +from onnx import AttributeProto, version_converter + +from ...onnx import ONNXDomain, check_onnx_model, get_captured_tensor_names +from ...pattern import PatternMatcher, PatternRewriter +from .base import BasePipe, OptimizationError, PipeConfig +from .cgir_rewrite_rules import ( + CGIR_REWRITE_CAPABILITIES, + CGIR_REWRITE_RULES, + CGIRModelRewriteRule, + CGIRRewriteRule, +) + + +if TYPE_CHECKING: + from onnx import ModelProto + + +logger = logging.getLogger(__name__) + + +class _CGIRPatternRewriter(PatternRewriter): + """Keep subgraph captures alive during CGIR-specific constant cleanup.""" + + def _remove_unused_constants(self, model: ModelProto) -> None: + graph = model.graph + consumed = {name for node in graph.node for name in node.input if name} + consumed.update(output.name for output in graph.output) + for node in graph.node: + for attribute in node.attribute: + if attribute.type == AttributeProto.GRAPH: + consumed.update(get_captured_tensor_names(attribute.g)) + elif attribute.type == AttributeProto.GRAPHS: + for subgraph in attribute.graphs: + consumed.update(get_captured_tensor_names(subgraph)) + + for node in list(graph.node): + if node.op_type == "Constant" and not consumed.intersection(node.output): + graph.node.remove(node) + removed_initializers = { + initializer.name for initializer in graph.initializer + if initializer.name not in consumed + } + for initializer in list(graph.initializer): + if initializer.name in removed_initializers: + graph.initializer.remove(initializer) + for graph_input in list(graph.input): + if graph_input.name in removed_initializers: + graph.input.remove(graph_input) + + +@dataclass +class CGIRRewritePipeConfig(PipeConfig): + """Configuration for target-specific CGIR rewrites.""" + + rules: list[CGIRRewriteRule | CGIRModelRewriteRule] = field(default_factory=list) + + +class CGIRRewritePipe(BasePipe[CGIRRewritePipeConfig]): + """Apply CGIR compatibility rewrites that may prepare the full model.""" + + name: ClassVar[str] = "cgir_rewrite" + capabilities: ClassVar[dict[str, Any]] = CGIR_REWRITE_CAPABILITIES + + @classmethod + def get_compatibility_options(cls) -> dict[str, bool]: + """Return options enabling every compatibility rule, without redundant aliases.""" + return {rule.capability.python_name: True for rule in CGIR_REWRITE_RULES} + + @classmethod + def build_config(cls, **kwargs: Any) -> CGIRRewritePipeConfig: + """Build the enabled CGIR rewrite configuration.""" + rules = [ + rule + for rule in CGIR_REWRITE_RULES + if any( + kwargs.get(capability.python_name) is True + for capability in (rule.capability, *rule.aliases) + ) + ] + return CGIRRewritePipeConfig(rules=rules) + + @classmethod + def should_process(cls, config: CGIRRewritePipeConfig) -> bool: + """Return whether at least one CGIR rewrite is enabled.""" + return bool(config.rules) + + def process( + self, + model: ModelProto, + config: CGIRRewritePipeConfig, + ) -> ModelProto: + """Apply each enabled rule once, matching against the preceding rule's result.""" + if not config.rules: + return model + + try: + rewritten_model = model + matcher: PatternMatcher | None = None + for rule in config.rules: + if isinstance(rule, CGIRModelRewriteRule): + if rule.minimum_opset and not any( + entry.domain == "" and entry.version >= rule.minimum_opset + for entry in rewritten_model.opset_import + ): + continue + prepared_model = rule.transform(rewritten_model) + if prepared_model is not rewritten_model: + logger.info( + "CGIR compatibility: %s applied model rewrite", + rule.capability.name, + ) + rewritten_model = prepared_model + matcher = None + continue + if matcher is None: + matcher = PatternMatcher(rewritten_model) + matcher.patterns.clear() + matcher.register_pattern(rule.source()) + matches = matcher.match() + if not matches: + continue + + current_opset = matcher.domain_versions.get(ONNXDomain.AI_ONNX) + if current_opset is None: + raise ValueError("Model does not declare the default ONNX opset") + if current_opset < rule.minimum_opset: + rewritten_model = version_converter.convert_version( + rewritten_model, + rule.minimum_opset, + ) + matcher = PatternMatcher(rewritten_model) + matcher.register_pattern(rule.source()) + matches = matcher.match() + if not matches: + continue + + original_outputs = {match.skeleton_match_result.output for match in matches} + rewritten_model = _CGIRPatternRewriter(rewritten_model).rewrite( + [(matches, rule.target)], + ) + residual_matcher = PatternMatcher(rewritten_model) + residual_matcher.register_pattern(rule.source()) + # A folded Cast can expose a downstream match; do not fold to a fixed point. + if any( + match.skeleton_match_result.output in original_outputs + for match in residual_matcher.match() + ): + raise ValueError( + f"CGIR rule {rule.capability.name} left selected source patterns " + "in the model", + ) + matcher = residual_matcher + if rule.warning: + logger.warning("%s (%d match(es))", rule.warning, len(matches)) + logger.info( + "CGIR compatibility: %s rewrote %d node(s)", rule.capability.name, len(matches) + ) + if rewritten_model is not model: + check_onnx_model(rewritten_model) + return rewritten_model + except OptimizationError: + raise + except Exception as error: + raise OptimizationError( + f"CGIR rewrite failed: {error}", + pipe_name=self.name, + cause=error, + ) from error + + +__all__ = ["CGIRRewritePipe", "CGIRRewritePipeConfig"] diff --git a/src/winml/modelkit/optim/pipes/cgir_rewrite_rules.py b/src/winml/modelkit/optim/pipes/cgir_rewrite_rules.py new file mode 100644 index 000000000..6e5e951ef --- /dev/null +++ b/src/winml/modelkit/optim/pipes/cgir_rewrite_rules.py @@ -0,0 +1,252 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Rule declarations for CGIR-specific model rewrites.""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass + +from onnx import ModelProto + +from ...pattern import Pattern +from ...pattern.cgc import ( + DFTWithStaticParametersPattern, + ExpandedPReluPattern, + GatherLinearGridSamplePattern, + GatherNDWithIdentityIndicesPattern, + LinearGridSamplePattern, + MatMulDFTPattern, + PReluWithFiniteSlopePattern, + ReshapedGatherNDPattern, + ResizeWithAsymmetricCoordinatesPattern, + ResizeWithCubicInterpolationPattern, + ResizeWithEmptyOptionalInputsPattern, + ResizeWithLinearInterpolationPattern, + ResizeWithOmittedOptionalInputsPattern, + ResizeWithTfHalfPixelForNNPattern, + cgc_constant_folding, + deduplicate_opset_imports, + eliminate_identity, + normalize_int32_dq, +) +from ..registry import BoolCapability, CapabilityCategory + + +@dataclass(frozen=True) +class CGIRRewriteRule: + """A CGIR rewrite and its model-level opset requirement.""" + + capability: BoolCapability + source: type[Pattern] + target: type[Pattern] + minimum_opset: int + aliases: tuple[BoolCapability, ...] = () + warning: str | None = None + + +@dataclass(frozen=True) +class CGIRModelRewriteRule: + """A model-level rewrite, independent of operator pattern matching.""" + + capability: BoolCapability + transform: Callable[[ModelProto], ModelProto] + aliases: tuple[BoolCapability, ...] = () + minimum_opset: int = 0 + + +DEDUPLICATE_OPSET_IMPORTS = BoolCapability( + name="deduplicate-opset-imports", + ort_name=None, + description="Remove repeated identical model opset imports before CGIR conversion", + category=CapabilityCategory.REWRITE, + default=False, +) + +OMIT_EMPTY_RESIZE_INPUTS = BoolCapability( + name="omit-empty-resize-inputs", + ort_name=None, + description="Represent empty Resize ROI and scales as omitted inputs for CGIR", + category=CapabilityCategory.REWRITE, + default=False, +) + +ELIMINATE_IDENTITY = BoolCapability( + name="eliminate-identity", + ort_name=None, + description="Eliminate internal tensor Identity aliases without changing graph IO for CGIR", + category=CapabilityCategory.REWRITE, + default=False, +) + +FOLD_CONSTANT_PAD_PADS = BoolCapability( + name="fold-constant-pad-pads", + ort_name=None, + description="Alias for cgc-constant-folding", + category=CapabilityCategory.REWRITE, + default=False, +) + +CGC_CONSTANT_FOLDING = BoolCapability( + name="cgc-constant-folding", + ort_name=None, + description="Fill FoundryToolbox folding gaps for Pad parameters and static shape subgraphs", + category=CapabilityCategory.REWRITE, + default=False, +) + +RESIZE_TF_HALF_PIXEL_FOR_NN_TO_ASYMMETRIC = BoolCapability( + name="resize-tf-half-pixel-for-nn-to-asymmetric", + ort_name=None, + description=( + "Use asymmetric coordinates for nearest/floor Resize with static positive integer scales" + ), + category=CapabilityCategory.REWRITE, + default=False, +) + +APPROXIMATE_CUBIC_RESIZE_WITH_LINEAR = BoolCapability( + name="approximate-cubic-resize-with-linear", + ort_name=None, + description=( + "Lossy cubic-to-linear Resize approximation without antialiasing or outside exclusion" + ), + category=CapabilityCategory.REWRITE, + default=False, +) + +GATHERND_TO_RESHAPE = BoolCapability( + name="gathernd-to-reshape", + ort_name=None, + description=( + "Replace unequal-rank GatherND with Reshape when static indices preserve all data in order" + ), + category=CapabilityCategory.REWRITE, + default=False, +) + +PRELU_TO_RELU = BoolCapability( + name="prelu-to-relu", + ort_name=None, + description="Decompose floating-point PRelu with finite constant slope into Relu/Neg/Mul/Sub", + category=CapabilityCategory.REWRITE, + default=False, +) + +DFT_TO_MATMUL = BoolCapability( + name="dft-to-matmul", + ort_name=None, + description=( + "Decompose DFT with static axis and signal length into real-valued matrix multiplications" + ), + category=CapabilityCategory.REWRITE, + default=False, +) + +GRIDSAMPLE_TO_GATHER = BoolCapability( + name="gridsample-to-gather", + ort_name=None, + description="Decompose 2D linear zero-padded GridSample into GatherND and FP32 interpolation", + category=CapabilityCategory.REWRITE, + default=False, +) + +NORMALIZE_INT32_DQ = BoolCapability( + name="normalize-int32-dq", + ort_name=None, + description="Omit constant INT32 DQ zero points and scalarize singleton scales for CGIR", + category=CapabilityCategory.REWRITE, + default=False, +) + +CGIR_REWRITE_RULES = ( + CGIRModelRewriteRule( + capability=NORMALIZE_INT32_DQ, + transform=normalize_int32_dq, + ), + CGIRModelRewriteRule( + capability=CGC_CONSTANT_FOLDING, + transform=cgc_constant_folding, + aliases=(FOLD_CONSTANT_PAD_PADS,), + ), + CGIRModelRewriteRule( + capability=DEDUPLICATE_OPSET_IMPORTS, + transform=deduplicate_opset_imports, + ), + CGIRRewriteRule( + capability=OMIT_EMPTY_RESIZE_INPUTS, + source=ResizeWithEmptyOptionalInputsPattern, + target=ResizeWithOmittedOptionalInputsPattern, + minimum_opset=13, + ), + CGIRRewriteRule( + capability=RESIZE_TF_HALF_PIXEL_FOR_NN_TO_ASYMMETRIC, + source=ResizeWithTfHalfPixelForNNPattern, + target=ResizeWithAsymmetricCoordinatesPattern, + minimum_opset=11, + ), + CGIRRewriteRule( + capability=APPROXIMATE_CUBIC_RESIZE_WITH_LINEAR, + source=ResizeWithCubicInterpolationPattern, + target=ResizeWithLinearInterpolationPattern, + minimum_opset=11, + warning=( + "Replacing cubic Resize with linear is a lossy approximation and may reduce accuracy." + ), + ), + CGIRRewriteRule( + capability=GATHERND_TO_RESHAPE, + source=GatherNDWithIdentityIndicesPattern, + target=ReshapedGatherNDPattern, + minimum_opset=11, + ), + CGIRRewriteRule( + capability=PRELU_TO_RELU, + source=PReluWithFiniteSlopePattern, + target=ExpandedPReluPattern, + minimum_opset=7, + ), + CGIRRewriteRule( + capability=DFT_TO_MATMUL, + source=DFTWithStaticParametersPattern, + target=MatMulDFTPattern, + minimum_opset=17, + ), + CGIRModelRewriteRule( + capability=ELIMINATE_IDENTITY, + transform=eliminate_identity, + ), + CGIRRewriteRule( + capability=GRIDSAMPLE_TO_GATHER, + source=LinearGridSamplePattern, + target=GatherLinearGridSamplePattern, + minimum_opset=16, + ), +) + +CGIR_REWRITE_CAPABILITIES = { + capability.name: capability + for rule in CGIR_REWRITE_RULES + for capability in (rule.capability, *rule.aliases) +} + + +__all__ = [ + "APPROXIMATE_CUBIC_RESIZE_WITH_LINEAR", + "CGC_CONSTANT_FOLDING", + "CGIR_REWRITE_CAPABILITIES", + "CGIR_REWRITE_RULES", + "DEDUPLICATE_OPSET_IMPORTS", + "DFT_TO_MATMUL", + "ELIMINATE_IDENTITY", + "FOLD_CONSTANT_PAD_PADS", + "GATHERND_TO_RESHAPE", + "GRIDSAMPLE_TO_GATHER", + "OMIT_EMPTY_RESIZE_INPUTS", + "PRELU_TO_RELU", + "RESIZE_TF_HALF_PIXEL_FOR_NN_TO_ASYMMETRIC", + "CGIRModelRewriteRule", + "CGIRRewriteRule", +] diff --git a/src/winml/modelkit/optim/pipes/graph.py b/src/winml/modelkit/optim/pipes/graph.py index 1fcb36bfb..10ca36315 100644 --- a/src/winml/modelkit/optim/pipes/graph.py +++ b/src/winml/modelkit/optim/pipes/graph.py @@ -56,6 +56,7 @@ # We exclude default=True items (ConstantFolding, IdentityElimination, etc.) # because ORT already enables those at Level 2 - no need to configure them. GRAPH_CAPABILITIES: dict[str, Any] = caps_dict( + graph_caps.ORT_GRAPH_OPTIMIZATION, # GELU fusions (all default=False) gelu.GELU_FUSION, gelu.FAST_GELU_FUSION, @@ -345,6 +346,8 @@ def build_config(cls, **kwargs: Any) -> ORTGraphPipeConfig: explicitly_disabled: list[str] = [] for cap in cls.capabilities.values(): + if cap is graph_caps.ORT_GRAPH_OPTIMIZATION: + continue if isinstance(cap, BoolCapability): user_value = kwargs.get(cap.python_name) if user_value is True: @@ -362,6 +365,7 @@ def build_config(cls, **kwargs: Any) -> ORTGraphPipeConfig: cap.python_name for cap in cls.capabilities.values() if isinstance(cap, BoolCapability) + and cap is not graph_caps.ORT_GRAPH_OPTIMIZATION and cap.default and cap.python_name not in explicitly_disabled ] @@ -371,6 +375,8 @@ def build_config(cls, **kwargs: Any) -> ORTGraphPipeConfig: verbose=verbose, ep_device=kwargs.get("ep_device"), ) + if kwargs.get("ort_graph_optimization") is False: + config.optimization_level = 0 # Explicitly disable capabilities that user set to False # This handles default=True caps like constant_folding diff --git a/src/winml/modelkit/pattern/cgc/__init__.py b/src/winml/modelkit/pattern/cgc/__init__.py new file mode 100644 index 000000000..82eda35ad --- /dev/null +++ b/src/winml/modelkit/pattern/cgc/__init__.py @@ -0,0 +1,45 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Opt-in CGC compatibility patterns and model metadata rewrites.""" + +from .cgc_constant_folding import cgc_constant_folding, fold_constant_pad_pads +from .dft_patterns import DFTWithStaticParametersPattern, MatMulDFTPattern +from .dq_rewrites import normalize_int32_dq +from .gathernd_patterns import GatherNDWithIdentityIndicesPattern, ReshapedGatherNDPattern +from .gridsample_patterns import GatherLinearGridSamplePattern, LinearGridSamplePattern +from .identity_rewrites import eliminate_identity +from .opset_rewrites import deduplicate_opset_imports +from .prelu_patterns import ExpandedPReluPattern, PReluWithFiniteSlopePattern +from .resize_patterns import ( + ResizeWithAsymmetricCoordinatesPattern, + ResizeWithCubicInterpolationPattern, + ResizeWithEmptyOptionalInputsPattern, + ResizeWithLinearInterpolationPattern, + ResizeWithOmittedOptionalInputsPattern, + ResizeWithTfHalfPixelForNNPattern, +) + + +__all__ = [ + "DFTWithStaticParametersPattern", + "ExpandedPReluPattern", + "GatherLinearGridSamplePattern", + "GatherNDWithIdentityIndicesPattern", + "LinearGridSamplePattern", + "MatMulDFTPattern", + "PReluWithFiniteSlopePattern", + "ReshapedGatherNDPattern", + "ResizeWithAsymmetricCoordinatesPattern", + "ResizeWithCubicInterpolationPattern", + "ResizeWithEmptyOptionalInputsPattern", + "ResizeWithLinearInterpolationPattern", + "ResizeWithOmittedOptionalInputsPattern", + "ResizeWithTfHalfPixelForNNPattern", + "cgc_constant_folding", + "deduplicate_opset_imports", + "eliminate_identity", + "fold_constant_pad_pads", + "normalize_int32_dq", +] diff --git a/src/winml/modelkit/pattern/cgc/cgc_constant_folding.py b/src/winml/modelkit/pattern/cgc/cgc_constant_folding.py new file mode 100644 index 000000000..f4cf2c49a --- /dev/null +++ b/src/winml/modelkit/pattern/cgc/cgc_constant_folding.py @@ -0,0 +1,346 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- + +"""Bounded constant folding for CGC Pad parameters and static shape subgraphs. + +FoundryToolbox currently lacks some constant folding needed by ONNX lowering. +These rewrites fill that gap without an ORT Session or execution-provider graph +transformations. Runtime floating-point computations are left unchanged. +""" + +from __future__ import annotations + +import logging +import math +from collections import Counter, deque +from typing import cast + +import numpy as np +from onnx import ( + AttributeProto, + GraphProto, + ModelProto, + TensorProto, + ValueInfoProto, + helper, + numpy_helper, + shape_inference, +) +from onnx.reference import ReferenceEvaluator + + +logger = logging.getLogger(__name__) +_MAX_ELEMENTS = 65536 +_MAX_NODES = 128 +_MAX_CACHED_ELEMENTS = 1048576 +_OPERATORS = {"ConstantOfShape", "Concat", "Reshape", "Slice", "Transpose", "Cast"} +_SHAPE_OPERATORS = _OPERATORS | { + "Mod", "Add", "Sub", "Mul", "Div", "Squeeze", "Unsqueeze", "Gather", "Equal", "Where", +} +_BROADCAST_OPERATORS = {"Mod", "Add", "Sub", "Mul", "Div", "Equal", "Where"} +_INTEGER_TYPES = { + TensorProto.INT8, TensorProto.INT16, TensorProto.INT32, + TensorProto.INT64, TensorProto.UINT8, TensorProto.UINT16, + TensorProto.UINT32, TensorProto.UINT64, TensorProto.BOOL, +} + + +class _ConstantParameters: + def __init__(self, model: ModelProto, *, static_shapes: bool = False) -> None: + self.model = model + self.static_shapes = static_shapes + self.shapes = { + value.name: value.type.tensor_type.shape + for value in [*model.graph.input, *model.graph.value_info] + if value.type.tensor_type.HasField("shape") + } + self.producers = { + name: node for node in model.graph.node for name in node.output if name + } + self.initializers = {value.name: value for value in model.graph.initializer} + self.inputs = {value.name for value in model.graph.input} + self.values: dict[str, np.ndarray] = {} + self.evaluated: set[str] = set() + self.cached_elements = 0 + + def tensor(self, value: TensorProto) -> np.ndarray: + if value.data_type not in _INTEGER_TYPES or math.prod(value.dims) > _MAX_ELEMENTS: + raise ValueError("Not a bounded integer tensor") + return numpy_helper.to_array(value) + + def evaluate(self, name: str, visited: set[str], active: set[str]) -> np.ndarray: + if name in self.inputs or name in active: + raise ValueError("Runtime input or cyclic dependency") + visited.add(name) + if len(visited) > _MAX_NODES or len(active) >= _MAX_NODES: + raise ValueError("Constant dependency budget exceeded") + if name in self.values: + return self.values[name] + active.add(name) + try: + if name in self.initializers: + result = self.tensor(self.initializers[name]) + else: + node = self.producers[name] + if node.domain not in {"", "ai.onnx"} or len(node.output) != 1: + raise ValueError("Unsupported constant producer") + if node.op_type == "Constant": + if len(node.attribute) != 1 or node.attribute[0].name != "value": + raise ValueError("Only tensor-valued Constants are supported") + result = self.tensor(node.attribute[0].t) + elif self.static_shapes and node.op_type == "Shape": + if len(node.input) != 1: + raise ValueError("Invalid Shape inputs") + shape = self.shapes.get(node.input[0]) + if shape is None: + tensor = self.initializers.get(node.input[0]) + if tensor is None: + raise ValueError("Unknown tensor rank") + shape = helper.make_tensor_type_proto( + tensor.data_type, list(tensor.dims) + ).tensor_type.shape + attributes = {attr.name: helper.get_attribute_value(attr) + for attr in node.attribute} + if set(attributes) - {"start", "end"}: + raise ValueError("Unsupported Shape attributes") + dimensions = list(shape.dim)[attributes.get("start", 0):attributes.get("end")] + if not all(dim.HasField("dim_value") and dim.dim_value >= 0 + for dim in dimensions): + raise ValueError("Shape contains dynamic dimensions") + result = np.asarray([dim.dim_value for dim in dimensions], dtype=np.int64) + else: + operators = _SHAPE_OPERATORS if self.static_shapes else _OPERATORS + if node.op_type not in operators: + raise ValueError("Unsupported integer constant operation") + inputs = { + item: self.evaluate(item, visited, active) for item in node.input if item + } + if sum(inputs[item].size for item in node.input if item) > _MAX_ELEMENTS: + raise ValueError("Constant operation input budget exceeded") + if node.op_type in _BROADCAST_OPERATORS: + output_shape = np.broadcast_shapes( + *(value.shape for value in inputs.values()), + ) + if math.prod(output_shape) > _MAX_ELEMENTS: + raise ValueError("Broadcast allocation budget exceeded") + if node.op_type == "Gather": + axis = next((attr.i for attr in node.attribute if attr.name == "axis"), 0) + data, indices = (inputs[item] for item in node.input) + axis %= data.ndim + output_shape = (*data.shape[:axis], *indices.shape, *data.shape[axis + 1:]) + if math.prod(output_shape) > _MAX_ELEMENTS: + raise ValueError("Gather allocation budget exceeded") + if node.op_type == "ConstantOfShape": + target_shape = inputs[node.input[0]] + if ( + target_shape.ndim != 1 or target_shape.dtype != np.int64 + or np.any(target_shape < 0) + or math.prod(int(dim) for dim in target_shape) > _MAX_ELEMENTS + or len(target_shape) > 32 + ): + raise ValueError("ConstantOfShape allocation budget exceeded") + if not node.attribute: + raise ValueError("Default ConstantOfShape output is floating point") + for attribute in node.attribute: + if attribute.name != "value": + raise ValueError("Unsupported ConstantOfShape attribute") + self.tensor(attribute.t) + if node.op_type == "Cast": + target = next(attr.i for attr in node.attribute if attr.name == "to") + if target not in _INTEGER_TYPES: + raise ValueError("Only integer Cast targets are supported") + fragment = helper.make_model( + helper.make_graph( + [node], "constant_pad_parameter", [], + [ValueInfoProto(name=name)], + [numpy_helper.from_array(value, key) for key, value in inputs.items()], + ), + opset_imports=list(self.model.opset_import), + ir_version=self.model.ir_version, + ) + result = cast("list[np.ndarray]", ReferenceEvaluator(fragment).run(None, {}))[0] + self.evaluated.add(name) + if result.dtype.kind not in "iub" or result.size > _MAX_ELEMENTS: + raise ValueError("Constant result exceeds supported type or size") + if self.cached_elements + result.size > _MAX_CACHED_ELEMENTS: + raise ValueError("Constant cache budget exceeded") + self.cached_elements += result.size + self.values[name] = result + return result + finally: + active.remove(name) + + +def _referenced_names(graph: GraphProto) -> list[str]: + names = [value.name for value in graph.output] + for annotation in graph.quantization_annotation: + names.append(annotation.tensor_name) + names.extend(item.value for item in annotation.quant_parameter_tensor_names) + for node in graph.node: + names.extend(name for name in node.input if name) + for attribute in node.attribute: + if attribute.type == AttributeProto.GRAPH: + names.extend(_referenced_names(attribute.g)) + elif attribute.type == AttributeProto.GRAPHS: + for child in attribute.graphs: + names.extend(_referenced_names(child)) + return names + + +def fold_constant_pad_pads(model: ModelProto) -> ModelProto: + """Fold constant integer Pad widths without specializing runtime input shapes. + + Only the main graph is rewritten. Nested graph captures and quantization + annotations conservatively protect shared producers from removal. + """ + versions = [item.version for item in model.opset_import if item.domain in {"", "ai.onnx"}] + if not versions or len(set(versions)) != 1 or versions[0] < 11: + return model + candidates = [ + (index, node) for index, node in enumerate(model.graph.node) + if node.domain in {"", "ai.onnx"} and node.op_type == "Pad" + and len(node.input) >= 2 and node.input[1] + ] + if not candidates: + return model + evaluator = _ConstantParameters(model) + types = { + value.name: value.type for value in + [*model.graph.input, *model.graph.value_info, *model.graph.output] + } + replacements: dict[int, np.ndarray] = {} + selected_dependencies: set[str] = set() + for index, node in candidates: + producer = evaluator.producers.get(node.input[1]) + if producer is None or producer.op_type == "Constant": + continue + try: + visited: set[str] = set() + pads = evaluator.evaluate(node.input[1], visited, set()) + if pads.dtype != np.int64 or pads.ndim != 1 or pads.size % 2: + continue + tensor_type = types.get(node.input[0]) + rank = ( + len(tensor_type.tensor_type.shape.dim) + if tensor_type is not None and tensor_type.tensor_type.HasField("shape") else None + ) + if len(node.input) > 3 and node.input[3]: + if versions[0] < 18: + continue + axes = evaluator.evaluate(node.input[3], visited, set()) + if axes.ndim != 1 or axes.dtype not in {np.dtype("int32"), np.dtype("int64")}: + continue + if pads.size != 2 * axes.size: + continue + if rank is not None: + if np.any(axes < -rank) or np.any(axes >= rank): + continue + if len({int(axis) % rank for axis in axes}) != axes.size: + continue + elif rank is not None and pads.size != 2 * rank: + continue + replacements[index] = pads + selected_dependencies.update(visited) + except ( + ValueError, KeyError, TypeError, IndexError, StopIteration, + NotImplementedError, OverflowError, + ): + logger.debug("Pad constant parameter is not foldable: %s", node.name, exc_info=True) + if not replacements: + return model + + rewritten = ModelProto() + rewritten.CopyFrom(model) + used_names = set(evaluator.producers) | set(evaluator.initializers) | evaluator.inputs + used_names.update(_referenced_names(model.graph)) + used_names.update(value.name for value in model.graph.value_info) + constants = [] + folded_names: dict[str, str] = {} + for index, value in replacements.items(): + node = rewritten.graph.node[index] + source = node.input[1] + if source not in folded_names: + name = source + "_folded_pads" + while name in used_names: + name += "_" + used_names.add(name) + folded_names[source] = name + constants.append(helper.make_node( + "Constant", [], [name], value=numpy_helper.from_array(value), + )) + node.input[1] = folded_names[source] + + references = Counter(_referenced_names(rewritten.graph)) + removable = { + node.output[0]: node for node in rewritten.graph.node + if len(node.output) == 1 + and node.output[0] in evaluator.evaluated & selected_dependencies + } + pending = deque(name for name in removable if not references[name]) + removed: set[str] = set() + while pending: + name = pending.popleft() + if name in removed: + continue + removed.add(name) + for source in removable[name].input: + references[source] -= 1 + if source in removable and not references[source]: + pending.append(source) + nodes = [node for node in rewritten.graph.node if not any(n in removed for n in node.output)] + del rewritten.graph.node[:] + rewritten.graph.node.extend([*constants, *nodes]) + infos = [value for value in rewritten.graph.value_info if value.name not in removed] + del rewritten.graph.value_info[:] + rewritten.graph.value_info.extend(infos) + logger.info("Folded constant pads for %d Pad node(s)", len(replacements)) + return rewritten + + +def cgc_constant_folding(model: ModelProto) -> ModelProto: + """Fill FoundryToolbox constant-folding gaps for Pad and static Shape chains. + + Only the main graph is changed. Pad widths use the existing bounded folder; + graphs containing Shape also fold bounded integer/boolean constant chains + to a fixed point with shape inference. This never freezes symbolic input + dimensions: callers must specialize inputs explicitly before this rule when + required. Casts of runtime data, including floating-point outputs, remain. + """ + prepared = fold_constant_pad_pads(model) + if not any(node.op_type == "Shape" and node.domain in {"", "ai.onnx"} + for node in prepared.graph.node): + return prepared + versions = {item.version for item in prepared.opset_import if item.domain in {"", "ai.onnx"}} + if len(versions) != 1 or next(iter(versions)) < 11: + return prepared + rewritten = ModelProto() + rewritten.CopyFrom(prepared) + changed = False + for _iteration in range(32): + for value_info in rewritten.graph.value_info: + value_info.type.tensor_type.ClearField("shape") + rewritten = shape_inference.infer_shapes(rewritten, strict_mode=True, data_prop=True) + evaluator = _ConstantParameters(rewritten, static_shapes=True) + folded = 0 + for node in rewritten.graph.node: + if node.op_type == "Constant" or len(node.output) != 1: + continue + try: + value = evaluator.evaluate(node.output[0], set(), set()) + except ( + ValueError, KeyError, TypeError, IndexError, StopIteration, + NotImplementedError, OverflowError, ZeroDivisionError, + ): + continue + node.CopyFrom(helper.make_node( + "Constant", [], list(node.output), value=numpy_helper.from_array(value), + )) + folded += 1 + if not folded: + return rewritten if changed else prepared + changed = True + logger.info("CGC constant folding: folded %d shape/integer node(s)", folded) + logger.warning("CGC constant folding reached the 32-round limit; retaining partial folding") + return rewritten diff --git a/src/winml/modelkit/pattern/cgc/dft_patterns.py b/src/winml/modelkit/pattern/cgc/dft_patterns.py new file mode 100644 index 000000000..f53f3676f --- /dev/null +++ b/src/winml/modelkit/pattern/cgc/dft_patterns.py @@ -0,0 +1,275 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Static-axis DFT decomposition into real-valued matrix multiplications.""" + +from __future__ import annotations + +from dataclasses import replace +from typing import TYPE_CHECKING, Any + +import numpy as np +from onnx import ModelProto, helper, numpy_helper +from onnx.defs import get_schema + +from ...onnx import ONNXDomain, SupportedONNXType +from .. import ( + InputInfo, + Pattern, + PatternMatchResult, + PatternSchema, + Skeleton, + opschema_to_pattern_schema, +) +from .utils import _depends_on_overridable_initializer, _static_tensor + + +if TYPE_CHECKING: + from .. import PatternMatcher, SkeletonMatchResult + + +_ONNX_DFT_SCHEMA = opschema_to_pattern_schema(get_schema("DFT", 20)) +_DFT_SCHEMA = replace(_ONNX_DFT_SCHEMA, inputs=_ONNX_DFT_SCHEMA.inputs[:1]) +_DFT_TYPES = set(get_schema("MatMul", 9).type_constraints[0].allowed_type_strs) & set( + get_schema("DFT", 20).type_constraints[0].allowed_type_strs +) + + +def _static_integer(name: str, matcher: PatternMatcher) -> int | None: + value = _static_tensor(name, matcher) + if ( + value is None + or value.ndim != 0 + or value.dtype not in (np.dtype(np.int32), np.dtype(np.int64)) + ): + return None + return int(value) + + +class _DFTPattern(Pattern): + def get_schema(self) -> PatternSchema: + return _DFT_SCHEMA + + def get_skeleton(self) -> Skeleton: + return Skeleton( + node_op_types=["DFT"], + node_domains=[ONNXDomain.AI_ONNX], + edges=[(-1, 0, 0, 0)], + exit_nodes=[0], + n_inputs=1, + ) + + def get_internal_constants_and_attributes( + self, + inputs: dict[str, np.ndarray], + attributes: dict[str, Any], + is_constant_map: dict[str, bool], + domain_versions: dict[ONNXDomain, int], + ) -> tuple[list[tuple[int, int, np.ndarray]], dict[tuple[int, str], Any]]: + return [], {} + + +class DFTWithStaticParametersPattern(_DFTPattern): + """Match DFT with known signal length, component count, and transform axis.""" + + def check_skeleton_result( + self, + skeleton_match_result: SkeletonMatchResult, + ) -> PatternMatchResult | None: + """Require static transform parameters without specializing batch dimensions.""" + node = skeleton_match_result.matched_nodes[0] + matcher = skeleton_match_result.matcher + opset = matcher.domain_versions.get(ONNXDomain.AI_ONNX, 0) + if opset < 17 or not node.input or len(node.output) != 1: + return None + input_name = node.input[0] + shape = matcher.get_tensor_shape(input_name) + dtype = matcher.get_tensor_type_str(input_name) + if ( + shape is None + or len(shape) < 2 + or shape[-1] not in (1, 2) + or dtype is None + or dtype not in _DFT_TYPES + or _depends_on_overridable_initializer(input_name, matcher) + ): + return None + attributes = {attr.name: helper.get_attribute_value(attr) for attr in node.attribute} + allowed_attributes = {"inverse", "onesided"} | ({"axis"} if opset < 20 else set()) + if attributes.keys() - allowed_attributes: + return None + inverse, onesided = attributes.get("inverse", 0), attributes.get("onesided", 0) + if inverse not in (0, 1) or onesided not in (0, 1): + return None + if onesided and (inverse or shape[-1] != 1): + return None + if opset < 20: + if len(node.input) > 2: + return None + axis = attributes.get("axis", 1) + else: + if len(node.input) > 3: + return None + axis = ( + _static_integer(node.input[2], matcher) + if len(node.input) > 2 and node.input[2] + else -2 + ) + if not isinstance(axis, int) or not -len(shape) <= axis < len(shape) - 1: + return None + axis %= len(shape) + if axis == len(shape) - 1: + return None + input_length = shape[axis] + if not isinstance(input_length, int) or input_length <= 0: + return None + length = ( + _static_integer(node.input[1], matcher) + if len(node.input) > 1 and node.input[1] + else input_length + ) + if length is None or length <= 0: + return None + return PatternMatchResult( + skeleton_match_result=skeleton_match_result, + schema_input_to_value={"input": input_name}, + schema_output_to_value={"output": node.output[0]}, + type_param_to_type={"T1": dtype}, + attributes={ + "_shape": shape, + "_axis": axis, + "_length": length, + "_inverse": inverse, + "_onesided": onesided, + "_ir_version": matcher.model.ir_version, + }, + input_infos={"input": InputInfo(name="input")}, + ) + + +class MatMulDFTPattern(_DFTPattern): + """Generate the DFT definition using sine/cosine bases and complex arithmetic.""" + + def get_onnx_model( + self, + inputs: dict[str, np.ndarray], + attributes: dict[str, Any], + is_constant_map: dict[str, bool], + output_dtypes: list[str], + domain_versions: dict[ONNXDomain, int], + prefix: str = "", + input_names: list[str] | None = None, + output_names: list[str] | None = None, + ) -> ModelProto: + """Preserve transform direction, normalization, truncation, and output layout.""" + del inputs, is_constant_map + if input_names is None or len(input_names) != 1: + raise ValueError("DFT rewrite requires one signal input") + if output_names is None or len(output_names) != 1: + raise ValueError("DFT rewrite requires one output") + shape = attributes["_shape"] + axis, length = attributes["_axis"], attributes["_length"] + inverse, onesided = attributes["_inverse"], attributes["_onesided"] + input_length = min(shape[axis], length) + output_length = length // 2 + 1 if onesided else length + output_type = SupportedONNXType.from_onnx_type(output_dtypes[0]) + nodes, initializers = [], [] + + def constant(name: str, value: np.ndarray) -> str: + name = f"{prefix}{name}" + initializers.append(numpy_helper.from_array(value, name)) + return name + + def op(op_type: str, names: list[str], suffix: str, **kwargs: Any) -> str: + result = f"{prefix}{suffix}" + nodes.append(helper.make_node(op_type, names, [result], name=result, **kwargs)) + return result + + # Only the retained input samples contribute. Missing samples are zero, + # so zero-padding does not need an activation-sized temporary tensor. + phase = ( + 2 + * np.pi + * ( + np.outer( + np.arange(input_length, dtype=np.float64), + np.arange(output_length, dtype=np.float64), + ) + % length + ) + / length + ) + normalization = 1 / length if inverse else 1 + cosine = constant("cosine", (np.cos(phase) * normalization).astype(output_type.np_type)) + sine = constant( + "sine", + (np.sin(phase) * (1 if inverse else -1) * normalization).astype(output_type.np_type), + ) + signal = input_names[0] + if shape[axis] > length: + signal = op( + "Slice", + [ + signal, + constant("starts", np.asarray([0], dtype=np.int64)), + constant("ends", np.asarray([length], dtype=np.int64)), + constant("slice_axes", np.asarray([axis], dtype=np.int64)), + ], + "truncated", + ) + permutation = [i for i in range(len(shape) - 1) if i != axis] + [axis, len(shape) - 1] + transposed = permutation != list(range(len(shape))) + if transposed: + signal = op("Transpose", [signal], "transposed", perm=permutation) + real = op( + "Gather", + [signal, constant("real_index", np.asarray(0, dtype=np.int64))], + "real", + axis=-1, + ) + real_out = op("MatMul", [real, cosine], "real_cosine") + imag_out = op("MatMul", [real, sine], "real_sine") + if shape[-1] == 2: + imag = op( + "Gather", + [signal, constant("imag_index", np.asarray(1, dtype=np.int64))], + "imag", + axis=-1, + ) + real_out = op("Sub", [real_out, op("MatMul", [imag, sine], "imag_sine")], "real_output") + imag_out = op( + "Add", [imag_out, op("MatMul", [imag, cosine], "imag_cosine")], "imag_output" + ) + component_axis = constant("component_axis", np.asarray([-1], dtype=np.int64)) + real_out = op("Unsqueeze", [real_out, component_axis], "real_component") + imag_out = op("Unsqueeze", [imag_out, component_axis], "imag_component") + result = op("Concat", [real_out, imag_out], "complex_output", axis=-1) + if transposed: + op("Transpose", [result], "restored", perm=np.argsort(permutation).tolist()) + nodes[-1].output[0] = output_names[0] + output_shape = list(shape) + output_shape[axis], output_shape[-1] = output_length, 2 + graph = helper.make_graph( + nodes, + f"{prefix}MatMulDFT", + [helper.make_tensor_value_info(input_names[0], output_type.tensor_proto_type, shape)], + [ + helper.make_tensor_value_info( + output_names[0], output_type.tensor_proto_type, output_shape + ) + ], + initializers, + ) + return helper.make_model( + graph, + producer_name="winmlcli-pattern-generator", + opset_imports=[ + helper.make_opsetid(domain.schema_domain, version) + for domain, version in domain_versions.items() + ], + ir_version=attributes["_ir_version"], + ) + + +__all__ = ["DFTWithStaticParametersPattern", "MatMulDFTPattern"] diff --git a/src/winml/modelkit/pattern/cgc/dq_rewrites.py b/src/winml/modelkit/pattern/cgc/dq_rewrites.py new file mode 100644 index 000000000..3c790668e --- /dev/null +++ b/src/winml/modelkit/pattern/cgc/dq_rewrites.py @@ -0,0 +1,82 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- + +"""Normalize constant INT32 DequantizeLinear parameters for CGC.""" + +import numpy as np +from onnx import AttributeProto, GraphProto, ModelProto, TensorProto, numpy_helper + + +def normalize_int32_dq(model: ModelProto) -> ModelProto: + """Omit immutable zero points and scalarize singleton INT32 DQ scales. + + Only local initializer-backed data and parameters are considered. Shared + initializers remain untouched; nested graphs are handled independently. + """ + versions = {opset.domain: opset.version for opset in model.opset_import} + result = ModelProto() + result.CopyFrom(model) + used_names: set[str] = set() + + def collect_names(graph: GraphProto) -> None: + used_names.update(value.name for value in graph.initializer) + used_names.update(value.name for value in [*graph.input, *graph.output, *graph.value_info]) + for node in graph.node: + used_names.update(node.input) + used_names.update(node.output) + for attribute in node.attribute: + if attribute.type == AttributeProto.GRAPH: + collect_names(attribute.g) + elif attribute.type == AttributeProto.GRAPHS: + for subgraph in attribute.graphs: + collect_names(subgraph) + + def rewrite(graph: GraphProto) -> bool: + changed = False + initializers = {value.name: value for value in graph.initializer} + inputs = {value.name for value in graph.input} + for node in graph.node: + for attribute in node.attribute: + if attribute.type == AttributeProto.GRAPH: + changed = rewrite(attribute.g) or changed + elif attribute.type == AttributeProto.GRAPHS: + for subgraph in attribute.graphs: + changed = rewrite(subgraph) or changed + supported = ( + node.domain == "" and versions.get("", 0) >= 10 + ) or (node.domain == "com.microsoft" and versions.get("com.microsoft") == 1) + if not supported or node.op_type != "DequantizeLinear" or len(node.input) not in (2, 3): + continue + if any(attribute.name != "axis" for attribute in node.attribute): + continue + data = initializers.get(node.input[0]) + scale = initializers.get(node.input[1]) + if data is None or data.data_type != TensorProto.INT32: + continue + if scale is None or scale.name in inputs or list(scale.dims) not in ([], [1]): + continue + zero_name = node.input[2] if len(node.input) == 3 else "" + if zero_name: + zero = initializers.get(zero_name) + if zero is None or zero.name in inputs or zero.data_type != TensorProto.INT32: + continue + if list(zero.dims) not in ([], [1]) or not np.all(numpy_helper.to_array(zero) == 0): + continue + if list(scale.dims) == [1]: + name = scale.name + "_cgc_scalar" + while name in used_names: + name += "_" + used_names.add(name) + scalar = numpy_helper.from_array(numpy_helper.to_array(scale).reshape(()), name) + graph.initializer.append(scalar) + node.input[1] = name + changed = True + if len(node.input) == 3: + del node.input[2] + changed = True + return changed + + collect_names(result.graph) + return result if rewrite(result.graph) else model diff --git a/src/winml/modelkit/pattern/cgc/gathernd_patterns.py b/src/winml/modelkit/pattern/cgc/gathernd_patterns.py new file mode 100644 index 000000000..8e208fb66 --- /dev/null +++ b/src/winml/modelkit/pattern/cgc/gathernd_patterns.py @@ -0,0 +1,162 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Equivalent GatherND-to-Reshape patterns for CGIR compatibility.""" + +from __future__ import annotations + +from math import prod +from typing import TYPE_CHECKING, Any + +import numpy as np +from onnx import ModelProto, helper, numpy_helper +from onnx.defs import get_schema + +from ...onnx import ONNXDomain, SupportedONNXType +from .. import InputInfo, PatternMatchResult, make_single_op_pattern +from .utils import _depends_on_overridable_initializer, _static_tensor + + +if TYPE_CHECKING: + from .. import SkeletonMatchResult + + +_GATHERND_SCHEMA, _SingleGatherNDPattern = make_single_op_pattern(get_schema("GatherND", 12)) + + +class GatherNDWithIdentityIndicesPattern(_SingleGatherNDPattern): # type: ignore[misc, valid-type] + """Match unequal-rank GatherND that visits every input slice in storage order.""" + + def check_skeleton_result( + self, + skeleton_match_result: SkeletonMatchResult, + ) -> PatternMatchResult | None: + """Require unequal ranks and complete, ordered indexing in every batch.""" + node = skeleton_match_result.matched_nodes[0] + matcher = skeleton_match_result.matcher + if len(node.input) != 2 or len(node.output) != 1: + return None + data_name, indices_name = node.input + indices = _static_tensor(indices_name, matcher) + data_shape = matcher.get_tensor_shape(data_name) + if ( + indices is None + or indices.dtype != np.dtype(np.int64) + or indices.ndim == 0 + or indices.size == 0 + or data_shape is None + or not data_shape + or any(not isinstance(dim, int) or dim <= 0 for dim in data_shape) + ): + return None + + if _depends_on_overridable_initializer(data_name, matcher): + return None + + batch_dims = next( + ( + helper.get_attribute_value(attr) + for attr in node.attribute + if attr.name == "batch_dims" + ), + 0, + ) + depth = indices.shape[-1] + if ( + not isinstance(batch_dims, int) + or not 0 <= batch_dims < min(len(data_shape), indices.ndim) + or not 1 <= depth <= len(data_shape) - batch_dims + or indices.shape[:batch_dims] != data_shape[:batch_dims] + ): + return None + indexed_shape = data_shape[batch_dims : batch_dims + depth] + slice_count = prod(indexed_shape) + if prod(indices.shape[batch_dims:-1]) != slice_count: + return None + output_shape = indices.shape[:-1] + data_shape[batch_dims + depth :] + if len(data_shape) == indices.ndim == len(output_shape): + return None + + # Compare coordinates, not data values: each batch must traverse all + # indexed slices once in row-major order. Negative equivalents are valid. + coordinates = indices.reshape(-1, slice_count, depth) + positions = np.arange(slice_count, dtype=np.int64) + stride = slice_count + for axis, extent in enumerate(indexed_shape): + stride //= extent + expected = (positions // stride) % extent + actual = coordinates[:, :, axis] + if not np.all((actual == expected) | (actual == expected - extent)): + return None + + type_mapping = self._infer_type_mapping(skeleton_match_result) + if "T" not in type_mapping: + return None + return PatternMatchResult( + skeleton_match_result=skeleton_match_result, + schema_input_to_value={"data": data_name, "indices": indices_name}, + schema_output_to_value={"output": node.output[0]}, + type_param_to_type=type_mapping, + attributes={ + "_data_shape": data_shape, + "_output_shape": output_shape, + "_ir_version": matcher.model.ir_version, + }, + input_infos={ + "data": InputInfo(name="data"), + "indices": InputInfo(name="indices"), + }, + ) + + +class ReshapedGatherNDPattern(_SingleGatherNDPattern): # type: ignore[misc, valid-type] + """Replace proven identity indexing with a statically shaped Reshape.""" + + def get_onnx_model( + self, + inputs: dict[str, np.ndarray], + attributes: dict[str, Any], + is_constant_map: dict[str, bool], + output_dtypes: list[str], + domain_versions: dict[ONNXDomain, int], + prefix: str = "", + input_names: list[str] | None = None, + output_names: list[str] | None = None, + ) -> ModelProto: + """Emit Reshape with a derived shape initializer and unchanged data type.""" + del inputs, is_constant_map + if input_names is None or len(input_names) != 2: + raise ValueError("GatherND rewrite requires data and indices names") + if output_names is None or len(output_names) != 1: + raise ValueError("GatherND rewrite requires one output name") + data_shape = attributes["_data_shape"] + output_shape = attributes["_output_shape"] + element_type = SupportedONNXType.from_onnx_type(output_dtypes[0]).tensor_proto_type + shape_name = f"{prefix}shape" + graph = helper.make_graph( + [ + helper.make_node( + "Reshape", + [input_names[0], shape_name], + output_names, + name=f"{prefix}Reshape", + ) + ], + f"{prefix}ReshapedGatherND", + [helper.make_tensor_value_info(input_names[0], element_type, data_shape)], + [helper.make_tensor_value_info(output_names[0], element_type, output_shape)], + [numpy_helper.from_array(np.asarray(output_shape, dtype=np.int64), shape_name)], + ) + return helper.make_model( + graph, + producer_name="winmlcli-pattern-generator", + opset_imports=[ + helper.make_opsetid(domain.schema_domain, version) + for domain, version in domain_versions.items() + ], + ir_version=attributes["_ir_version"], + ) + + +__all__ = ["GatherNDWithIdentityIndicesPattern", "ReshapedGatherNDPattern"] diff --git a/src/winml/modelkit/pattern/cgc/gridsample_patterns.py b/src/winml/modelkit/pattern/cgc/gridsample_patterns.py new file mode 100644 index 000000000..8b92c073c --- /dev/null +++ b/src/winml/modelkit/pattern/cgc/gridsample_patterns.py @@ -0,0 +1,204 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- + +"""Lower 2D linear, zero-padded GridSample to indexed interpolation.""" + +from __future__ import annotations + +from typing import Any + +import numpy as np +from onnx import ModelProto, TensorProto, helper, numpy_helper +from onnx.defs import get_schema + +from ...onnx import ONNXDomain +from .. import InputInfo, PatternMatchResult, SkeletonMatchResult, make_single_op_pattern + + +_SCHEMA, _GridSamplePattern = make_single_op_pattern(get_schema("GridSample", 16)) + + +class LinearGridSamplePattern(_GridSamplePattern): # type: ignore[misc, valid-type] + """Match floating-point 2D sampling with known spatial and channel dimensions.""" + + def check_skeleton_result( + self, skeleton_match_result: SkeletonMatchResult, + ) -> PatternMatchResult | None: + """Reject unsupported modes, types and spatial dimensions.""" + node = skeleton_match_result.matched_nodes[0] + matcher = skeleton_match_result.matcher + version = matcher.domain_versions.get(ONNXDomain.AI_ONNX, 0) + if node.domain not in {"", ONNXDomain.AI_ONNX.value} or version < 16: + return None + attributes = {attr.name: helper.get_attribute_value(attr) for attr in node.attribute} + if ( + set(attributes) - {"mode", "padding_mode", "align_corners"} + or attributes.get("mode", b"linear" if version >= 20 else b"bilinear") + != (b"linear" if version >= 20 else b"bilinear") + or attributes.get("padding_mode", b"zeros") != b"zeros" + or attributes.get("align_corners", 0) not in {0, 1} + or len(node.input) != 2 or len(node.output) != 1 + ): + return None + shapes = [matcher.get_tensor_shape(name) for name in node.input] + types = [matcher.get_tensor_type_str(name) for name in node.input] + if types[0] not in {"tensor(float)", "tensor(float16)"} or types[1] != types[0]: + return None + if any(shape is None or len(shape) != 4 for shape in shapes): + return None + assert shapes[0] is not None + assert shapes[1] is not None + if shapes[1][-1] != 2 or any( + not isinstance(dim, int) or dim <= 0 + for dim in [*shapes[0][1:], *shapes[1][1:3]] + ): + return None + if ( + isinstance(shapes[0][0], int) and isinstance(shapes[1][0], int) + and shapes[0][0] != shapes[1][0] + ): + return None + attributes.update({"_shapes": shapes, "_ir_version": matcher.model.ir_version}) + return PatternMatchResult( + skeleton_match_result=skeleton_match_result, + schema_input_to_value={"X": node.input[0], "grid": node.input[1]}, + schema_output_to_value={"Y": node.output[0]}, + type_param_to_type={"T1": types[0], "T2": types[1]}, + attributes=attributes, + input_infos={"X": InputInfo(name="X"), "grid": InputInfo(name="grid")}, + ) + + +class GatherLinearGridSamplePattern(_GridSamplePattern): # type: ignore[misc, valid-type] + """Interpolate four samples with explicit runtime batch coordinates. + + GatherND uses batch_dims=0 to avoid the symbolic shape inference defect in + https://github.com/microsoft/onnxruntime/pull/24206. + """ + + def get_onnx_model( + self, + inputs: dict[str, np.ndarray], + attributes: dict[str, Any], + is_constant_map: dict[str, bool], + output_dtypes: list[str], + domain_versions: dict[ONNXDomain, int], + prefix: str = "", + input_names: list[str] | None = None, + output_names: list[str] | None = None, + ) -> ModelProto: + """Build batched indexed interpolation without specializing batch.""" + del inputs, is_constant_map + data_shape, grid_shape = attributes["_shapes"] + _, channels, height, width = data_shape + align = attributes.get("align_corners", 0) + fp16 = output_dtypes[0] == "tensor(float16)" + element_type = TensorProto.FLOAT16 if fp16 else TensorProto.FLOAT + nodes = [] + + def constant(label: str, value: Any) -> str: + name = f"{prefix}{label}" + nodes.append(helper.make_node( + "Constant", [], [name], value=numpy_helper.from_array(np.asarray(value)), + )) + return name + + def operation(kind: str, *arguments: str, **attrs: Any) -> str: + name = f"{prefix}value_{len(nodes)}" + nodes.append(helper.make_node(kind, list(arguments), [name], **attrs)) + return name + + assert input_names is not None + assert output_names is not None + data, grid = input_names + if fp16: + data = operation("Cast", data, to=TensorProto.FLOAT) + grid = operation("Cast", grid, to=TensorProto.FLOAT) + zero = constant("zero", np.float32(0)) + one = constant("one", np.float32(1)) + half = constant("half", np.float32(0.5)) if not align else None + axes = constant("axes", np.asarray([-1], np.int64)) + coordinates = [] + for axis, size in enumerate((width, height)): + component = operation( + "Gather", grid, constant(f"axis_{axis}", np.int64(axis)), axis=3, + ) + scale = constant(f"scale_{axis}", np.float32((size - 1 if align else size) / 2)) + pixel = operation("Mul", operation("Add", component, one), scale) + if not align: + assert half is not None + pixel = operation("Sub", pixel, half) + lower = operation("Floor", pixel) + fraction = operation("Sub", pixel, lower) + coordinates.append((lower, operation("Add", lower, one), + operation("Sub", one, fraction), fraction)) + source = operation( + "Reshape", operation("Transpose", data, perm=[0, 2, 3, 1]), + constant("source_shape", np.asarray([0, -1, channels], np.int64)), + ) + stride = constant("stride", np.int64(width)) + index_shape = constant("index_shape", np.asarray([0, -1, 1], np.int64)) + sampled_shape = constant( + "sampled_shape", np.asarray([0, grid_shape[1], grid_shape[2], channels], np.int64), + ) + limits = [constant(f"limit_{axis}", np.float32(size - 1)) + for axis, size in enumerate((width, height))] + batch_indices = None + terms = [] + for row in range(2): + for column in range(2): + valid, indices = [], [] + for coord, limit in zip( + [coordinates[0][column], coordinates[1][row]], limits, strict=True, + ): + valid.append(operation( + "And", operation("GreaterOrEqual", coord, zero), + operation("LessOrEqual", coord, limit), + )) + indices.append(operation( + "Cast", operation("Clip", coord, zero, limit), to=TensorProto.INT64, + )) + linear = operation("Add", indices[0], operation("Mul", indices[1], stride)) + spatial_indices = operation("Reshape", linear, index_shape) + if batch_indices is None: + indices_shape = operation("Shape", spatial_indices) + batch_zero = constant("batch_zero", np.int64(0)) + batch_one = constant("batch_one", np.int64(1)) + batch_size = operation("Gather", indices_shape, batch_zero, axis=0) + batch_range = operation("Range", batch_zero, batch_size, batch_one) + batch_indices = operation( + "Reshape", batch_range, + constant("batch_shape", np.asarray([-1, 1, 1], np.int64)), + ) + batch_indices = operation("Expand", batch_indices, indices_shape) + sampled = operation( + "GatherND", source, + operation("Concat", batch_indices, spatial_indices, axis=2), batch_dims=0, + ) + sampled = operation("Reshape", sampled, sampled_shape) + sampled = operation( + "Where", operation("Unsqueeze", operation("And", *valid), axes), sampled, zero, + ) + weight = operation("Mul", coordinates[0][column + 2], coordinates[1][row + 2]) + terms.append(operation("Mul", sampled, operation("Unsqueeze", weight, axes))) + result = operation("Add", operation("Add", terms[0], terms[1]), + operation("Add", terms[2], terms[3])) + result = operation("Transpose", result, perm=[0, 3, 1, 2]) + if fp16: + operation("Cast", result, to=TensorProto.FLOAT16) + nodes[-1].output[0] = output_names[0] + output_shape = [data_shape[0], channels, grid_shape[1], grid_shape[2]] + return helper.make_model( + helper.make_graph(nodes, f"{prefix}GridSampleToGather", + [helper.make_tensor_value_info(name, element_type, shape) + for name, shape in zip(input_names, [data_shape, grid_shape], strict=True)], + [helper.make_tensor_value_info(output_names[0], element_type, output_shape)]), + opset_imports=[helper.make_opsetid(domain.schema_domain, version) + for domain, version in domain_versions.items()], + ir_version=attributes["_ir_version"], + ) + + +__all__ = ["GatherLinearGridSamplePattern", "LinearGridSamplePattern"] diff --git a/src/winml/modelkit/pattern/cgc/identity_rewrites.py b/src/winml/modelkit/pattern/cgc/identity_rewrites.py new file mode 100644 index 000000000..665f949ab --- /dev/null +++ b/src/winml/modelkit/pattern/cgc/identity_rewrites.py @@ -0,0 +1,206 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Eliminate internal tensor Identity aliases for IX compatibility.""" + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING + +from onnx import AttributeProto, GraphProto, ModelProto, NodeProto, TensorProto, TypeProto, helper + +from ...onnx import ONNXDomain + + +if TYPE_CHECKING: + from collections.abc import Iterator + + +logger = logging.getLogger(__name__) + + +def _subgraphs(node: NodeProto) -> Iterator[GraphProto]: + for attribute in node.attribute: + if attribute.type == AttributeProto.GRAPH: + yield attribute.g + elif attribute.type == AttributeProto.GRAPHS: + yield from attribute.graphs + + +def _local_names(graph: GraphProto) -> set[str]: + return ( + {value.name for value in graph.input} + | {value.name for value in graph.initializer} + | {value.values.name for value in graph.sparse_initializer} + | {name for node in graph.node for name in node.output if name} + ) + + +def _types(graph: GraphProto, outer: dict[str, TypeProto]) -> dict[str, TypeProto]: + local = _local_names(graph) + types = {name: value for name, value in outer.items() if name not in local} + for value in (*graph.value_info, *graph.input, *graph.output): + types[value.name] = value.type + for tensor in graph.initializer: + types.setdefault(tensor.name, helper.make_tensor_type_proto(tensor.data_type, tensor.dims)) + return types + + +def _compatible_tensor(source: TypeProto | None, output: TypeProto | None) -> bool: + if source is None or not source.HasField("tensor_type"): + return False + if output is None: + return True + if not output.HasField("tensor_type"): + return False + left, right = source.tensor_type, output.tensor_type + if left.elem_type and right.elem_type and left.elem_type != right.elem_type: + return False + if left.HasField("shape") and right.HasField("shape"): + if len(left.shape.dim) != len(right.shape.dim): + return False + for a, b in zip(left.shape.dim, right.shape.dim, strict=True): + if a.HasField("dim_value") and b.HasField("dim_value") and a.dim_value != b.dim_value: + return False + if a.dim_param and b.dim_param and a.dim_param != b.dim_param: + return False + return True + + +def _redirect_uses( + graph: GraphProto, old: str, new: str, *, apply: bool, shadowed: bool = False +) -> bool: + if any(value.name == old for value in graph.output): + return False + # Quantization annotations can name aliases independently of node inputs. + if any( + annotation.tensor_name == old + or any(parameter.value == old for parameter in annotation.quant_parameter_tensor_names) + for annotation in graph.quantization_annotation + ): + return False + for node in graph.node: + for index, name in enumerate(node.input): + if name == old: + if shadowed: + return False + if apply: + node.input[index] = new + for child in _subgraphs(node): + local = _local_names(child) + if old in local: + continue + if not _redirect_uses( + child, old, new, apply=apply, shadowed=shadowed or new in local + ): + return False + if apply: + retained = [value for value in graph.value_info if value.name != old] + del graph.value_info[:] + graph.value_info.extend(retained) + return True + + +def _eliminate(graph: GraphProto, outer: dict[str, TypeProto]) -> int: + types = _types(graph, outer) + removed = 0 + for node in list(graph.node): + if node.domain not in {"", ONNXDomain.AI_ONNX.value} or node.op_type != "Identity": + continue + if len(node.input) != 1 or len(node.output) != 1 or node.attribute: + continue + source, output = node.input[0], node.output[0] + if not source or not output or source == output: + continue + if not _compatible_tensor(types.get(source), types.get(output)): + logger.debug( + "Retaining Identity %r: tensor types are unknown or incompatible", node.name, + ) + continue + if not _redirect_uses(graph, output, source, apply=False): + logger.debug( + "Retaining Identity %r: output interface or scoped alias is protected", node.name + ) + continue + _redirect_uses(graph, output, source, apply=True) + graph.node.remove(node) + types.pop(output, None) + removed += 1 + for node in graph.node: + for child in _subgraphs(node): + removed += _eliminate(child, types) + return removed + + +def _reshape_output_aliases(model: ModelProto) -> int: + versions = [entry.version for entry in model.opset_import if entry.domain == ""] + graph = model.graph + if len(versions) != 1 or versions[0] < 5 or any(list(_subgraphs(node)) for node in graph.node): + return 0 + types: dict[str, TypeProto] = {} + for value in (*graph.input, *graph.output, *graph.value_info): + if value.name in types and types[value.name] != value.type: + return 0 + types[value.name] = value.type + outputs = {value.name for value in graph.output} + inputs = {value.name for value in graph.input} + protected = set() + for annotation in graph.quantization_annotation: + protected.add(annotation.tensor_name) + protected.update(parameter.value for parameter in annotation.quant_parameter_tensor_names) + names = _local_names(graph) | set(types) | protected + names.update(name for node in graph.node for name in node.input) + rewritten = 0 + for node in graph.node: + if node.domain or node.op_type != "Identity" or node.attribute: + continue + if len(node.input) != 1 or len(node.output) != 1: + continue + source, output = node.input[0], node.output[0] + if (not source or source == output or output not in outputs or output in inputs + or source in protected or output in protected): + continue + source_type, output_type = types.get(source), types.get(output) + if (source_type is None or source_type != output_type + or not source_type.HasField("tensor_type")): + continue + tensor = source_type.tensor_type + if tensor.elem_type != TensorProto.FLOAT or not tensor.HasField("shape"): + continue + dimensions = tensor.shape.dim + if not dimensions or any(not dim.HasField("dim_value") or dim.dim_value <= 0 + for dim in dimensions): + continue + shape_name = output + "_identity_shape" + while shape_name in names: + shape_name += "_" + names.add(shape_name) + graph.initializer.append(helper.make_tensor( + shape_name, TensorProto.INT64, [len(dimensions)], + [dim.dim_value for dim in dimensions], + )) + node.op_type = "Reshape" + node.input.append(shape_name) + rewritten += 1 + return rewritten + + +def eliminate_identity(model: ModelProto) -> ModelProto: + """Remove safe internal tensor Identities, preserving graph IO and lexical bindings. + + Top-level static positive-shape FP32 output aliases use Reshape in models + without subgraphs. Other protected aliases and lexical bindings are retained. + """ + result = ModelProto() + result.CopyFrom(model) + reshaped = _reshape_output_aliases(result) + removed = _eliminate(result.graph, {}) + if not removed and not reshaped: + return model + logger.info( + "CGIR compatibility: eliminate-identity removed %d node(s), reshaped %d output alias(es)", + removed, reshaped, + ) + return result diff --git a/src/winml/modelkit/pattern/cgc/opset_rewrites.py b/src/winml/modelkit/pattern/cgc/opset_rewrites.py new file mode 100644 index 000000000..130e5a24c --- /dev/null +++ b/src/winml/modelkit/pattern/cgc/opset_rewrites.py @@ -0,0 +1,30 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Model-level opset metadata rewrites for CGC compatibility.""" + +from onnx import ModelProto + + +def deduplicate_opset_imports(model: ModelProto) -> ModelProto: + """Remove identical model opset imports, rejecting conflicting versions.""" + versions: dict[str, int] = {} + for opset in model.opset_import: + if opset.domain in versions and versions[opset.domain] != opset.version: + raise ValueError( + f"Conflicting opset imports for domain {opset.domain!r}: " + f"{versions[opset.domain]} and {opset.version}" + ) + versions[opset.domain] = opset.version + if len(versions) == len(model.opset_import): + return model + result = ModelProto() + result.CopyFrom(model) + del result.opset_import[:] + seen: set[str] = set() + for opset in model.opset_import: + if opset.domain not in seen: + result.opset_import.append(opset) + seen.add(opset.domain) + return result diff --git a/src/winml/modelkit/pattern/cgc/prelu_patterns.py b/src/winml/modelkit/pattern/cgc/prelu_patterns.py new file mode 100644 index 000000000..4c67b9e32 --- /dev/null +++ b/src/winml/modelkit/pattern/cgc/prelu_patterns.py @@ -0,0 +1,132 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""PRelu decomposition for the missing IX ONNX legalization.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +import numpy as np +from onnx import ModelProto, helper +from onnx.defs import get_schema + +from ...onnx import ONNXDomain, SupportedONNXType +from .. import InputInfo, PatternMatchResult, make_single_op_pattern +from .utils import _static_tensor + + +if TYPE_CHECKING: + from .. import SkeletonMatchResult + + +_PRELU_SCHEMA, _SinglePReluPattern = make_single_op_pattern(get_schema("PRelu", 9)) +_RELU_TYPES = set(get_schema("Relu", 6).type_constraints[0].allowed_type_strs) + + +class PReluWithFiniteSlopePattern(_SinglePReluPattern): # type: ignore[misc, valid-type] + """Match floating-point PRelu with a finite, non-overridable constant slope.""" + + def check_skeleton_result( + self, + skeleton_match_result: SkeletonMatchResult, + ) -> PatternMatchResult | None: + """Preserve broadcasting and reject slopes that would introduce NaNs.""" + node = skeleton_match_result.matched_nodes[0] + matcher = skeleton_match_result.matcher + if ( + len(node.input) != 2 + or len(node.output) != 1 + or matcher.domain_versions.get(ONNXDomain.AI_ONNX, 0) < 7 + ): + return None + data_name, slope_name = node.input + data_type = matcher.get_tensor_type_str(data_name) + if ( + data_type is None + or data_type not in _RELU_TYPES + or matcher.get_tensor_type_str(slope_name) != data_type + ): + return None + slope = _static_tensor(slope_name, matcher) + # With non-finite slopes, even x > 0 would evaluate slope * 0 to NaN. + if slope is None or not np.all(np.isfinite(slope)): + return None + return PatternMatchResult( + skeleton_match_result=skeleton_match_result, + schema_input_to_value={"X": data_name, "slope": slope_name}, + schema_output_to_value={"Y": node.output[0]}, + type_param_to_type={"T": data_type}, + attributes={ + "_data_shape": matcher.get_tensor_shape(data_name), + "_slope_shape": slope.shape, + "_ir_version": matcher.model.ir_version, + }, + input_infos={ + "X": InputInfo(name="X"), + "slope": InputInfo(name="slope"), + }, + ) + + +class ExpandedPReluPattern(_SinglePReluPattern): # type: ignore[misc, valid-type] + """Generate Relu(x) - slope * Relu(-x) without folding or copying the slope.""" + + def get_onnx_model( + self, + inputs: dict[str, np.ndarray], + attributes: dict[str, Any], + is_constant_map: dict[str, bool], + output_dtypes: list[str], + domain_versions: dict[ONNXDomain, int], + prefix: str = "", + input_names: list[str] | None = None, + output_names: list[str] | None = None, + ) -> ModelProto: + """Build the floating-point decomposition with the original tensor boundaries.""" + del inputs, is_constant_map + if input_names is None or len(input_names) != 2: + raise ValueError("PRelu rewrite requires data and slope names") + if output_names is None or len(output_names) != 1: + raise ValueError("PRelu rewrite requires one output name") + data_name, slope_name = input_names + positive = f"{prefix}positive" + negated = f"{prefix}negated" + negative = f"{prefix}negative" + scaled = f"{prefix}scaled" + nodes = [ + helper.make_node("Relu", [data_name], [positive], name=f"{prefix}PositiveRelu"), + helper.make_node("Neg", [data_name], [negated], name=f"{prefix}Neg"), + helper.make_node("Relu", [negated], [negative], name=f"{prefix}NegativeRelu"), + helper.make_node("Mul", [slope_name, negative], [scaled], name=f"{prefix}Mul"), + helper.make_node("Sub", [positive, scaled], output_names, name=f"{prefix}Sub"), + ] + element_type = SupportedONNXType.from_onnx_type(output_dtypes[0]).tensor_proto_type + graph = helper.make_graph( + nodes, + f"{prefix}ExpandedPRelu", + [ + helper.make_tensor_value_info(data_name, element_type, attributes["_data_shape"]), + helper.make_tensor_value_info(slope_name, element_type, attributes["_slope_shape"]), + ], + [ + helper.make_tensor_value_info( + output_names[0], + element_type, + attributes["_data_shape"], + ) + ], + ) + return helper.make_model( + graph, + producer_name="winmlcli-pattern-generator", + opset_imports=[ + helper.make_opsetid(domain.schema_domain, version) + for domain, version in domain_versions.items() + ], + ir_version=attributes["_ir_version"], + ) + + +__all__ = ["ExpandedPReluPattern", "PReluWithFiniteSlopePattern"] diff --git a/src/winml/modelkit/pattern/cgc/resize_patterns.py b/src/winml/modelkit/pattern/cgc/resize_patterns.py new file mode 100644 index 000000000..f6e0232b6 --- /dev/null +++ b/src/winml/modelkit/pattern/cgc/resize_patterns.py @@ -0,0 +1,470 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- + +"""Explicit Resize compatibility rewrites for CGIR.""" + +from __future__ import annotations + +from dataclasses import replace +from typing import TYPE_CHECKING, Any + +import numpy as np +from onnx import ModelProto, TensorProto, helper +from onnx.defs import OpSchema, get_schema + +from ...onnx import ONNXDomain, SupportedONNXType +from .. import ( + InputInfo, + Pattern, + PatternMatchResult, + PatternSchema, + Skeleton, + opschema_to_pattern_schema, +) +from ..utils import get_attribute_proto_value + + +if TYPE_CHECKING: + from .. import PatternMatcher, SkeletonMatchResult + + +_ONNX_RESIZE_SCHEMA = get_schema("Resize", 13) +_RESIZE_SCHEMA: PatternSchema = opschema_to_pattern_schema(_ONNX_RESIZE_SCHEMA) +_RESIZE_OMITTED_INPUTS_SCHEMA = PatternSchema( + name=_RESIZE_SCHEMA.name, + doc=_RESIZE_SCHEMA.doc, + inputs=[ + _RESIZE_SCHEMA.inputs[0], + OpSchema.FormalParameter( + name="resize_parameter", + type_str="TResizeParameter", + description="The effective non-empty scales or sizes input.", + param_option=OpSchema.FormalParameterOption.Single, + is_homogeneous=True, + min_arity=1, + differentiation_category=OpSchema.DifferentiationCategory.Differentiable, + ), + ], + outputs=_RESIZE_SCHEMA.outputs, + type_constraints=[ + *_RESIZE_SCHEMA.type_constraints, + OpSchema.TypeConstraintParam( + type_param_str="TResizeParameter", + allowed_type_strs=["tensor(float)", "tensor(int64)"], + description="Constrain the effective Resize parameter.", + ), + ], + attributes=_RESIZE_SCHEMA.attributes, +) + + +def _is_empty_shape(shape: tuple[int | str | None, ...] | None) -> bool: + return shape is not None and any(dimension == 0 for dimension in shape) + + +def _is_empty_resize_input(name: str, matcher: PatternMatcher) -> bool: + """Require static emptiness without relying on an overridable input's default.""" + if not name or not _is_empty_shape(matcher.get_tensor_shape(name)): + return False + pending = [name] + visited: set[str] = set() + while pending: + current = pending.pop() + if not current or current in visited: + continue + visited.add(current) + producer = matcher.producer_lookup.get(current) + if producer is None or producer[2] == "GraphInput": + return False + node = matcher.node_lookup.get(producer[0]) + if node is not None: + pending.extend(node.input) + return True + + +class _ResizeOptionalInputsPattern(Pattern): + def get_skeleton(self) -> Skeleton: + return Skeleton( + node_op_types=["Resize"], + node_domains=[ONNXDomain.AI_ONNX], + edges=[ + (-1, 0, 0, 0), + ], + exit_nodes=[0], + n_inputs=1, + ) + + def get_schema(self) -> PatternSchema: + return _RESIZE_OMITTED_INPUTS_SCHEMA + + def _infer_schema_attributes( + self, + skeleton_match_result: SkeletonMatchResult, + ) -> dict[str, Any]: + node = skeleton_match_result.matched_nodes[0] + return { + attribute.name: get_attribute_proto_value( + attribute, + replace_float_with_dummy=False, + ) + for attribute in node.attribute + } + + def get_internal_constants_and_attributes( + self, + inputs: dict[str, np.ndarray], + attributes: dict[str, Any], + is_constant_map: dict[str, bool], + domain_versions: dict[ONNXDomain, int], + ) -> tuple[list[tuple[int, int, np.ndarray]], dict[tuple[int, str], Any]]: + return [], {(0, name): value for name, value in attributes.items()} + + +class ResizeWithEmptyOptionalInputsPattern(_ResizeOptionalInputsPattern): + """Match Resize inputs represented by named empty tensors.""" + + def check_skeleton_result( + self, + skeleton_match_result: SkeletonMatchResult, + ) -> PatternMatchResult | None: + """Omit statically empty ROI/scales while retaining the effective parameter.""" + node = skeleton_match_result.matched_nodes[0] + matcher = skeleton_match_result.matcher + if matcher.domain_versions.get(ONNXDomain.AI_ONNX, 0) < 11: + return None + attributes = self._infer_schema_attributes(skeleton_match_result) + if attributes.get("coordinate_transformation_mode") == "tf_crop_and_resize": + return None + + roi_name = node.input[1] if len(node.input) > 1 else "" + scales_name = node.input[2] if len(node.input) > 2 else "" + sizes_name = node.input[3] if len(node.input) > 3 else "" + + roi_is_empty = _is_empty_resize_input(roi_name, matcher) + if roi_name and not roi_is_empty: + return None + + scales_is_empty = _is_empty_resize_input(scales_name, matcher) + if sizes_name: + if scales_name and not scales_is_empty: + return None + resize_parameter_name = sizes_name + resize_parameter_slot = 3 + elif scales_name and not scales_is_empty: + resize_parameter_name = scales_name + resize_parameter_slot = 2 + else: + return None + + if not roi_is_empty and not scales_is_empty: + return None + + skeleton_match_result = replace( + skeleton_match_result, + inputs=[node.input[0], resize_parameter_name], + ) + type_param_to_type = self._infer_type_mapping(skeleton_match_result) + if "T1" not in type_param_to_type: + return None + + attributes["_resize_parameter_slot"] = resize_parameter_slot + attributes["_ir_version"] = matcher.model.ir_version + return PatternMatchResult( + skeleton_match_result=skeleton_match_result, + schema_input_to_value={ + "X": node.input[0], + "resize_parameter": resize_parameter_name, + }, + schema_output_to_value={"Y": node.output[0]}, + type_param_to_type=type_param_to_type, + attributes=attributes, + input_infos={ + "X": InputInfo(name="X"), + "resize_parameter": InputInfo(name="resize_parameter"), + }, + ) + + +class ResizeWithOmittedOptionalInputsPattern(_ResizeOptionalInputsPattern): + """Generate Resize with empty ROI and scales represented as omitted inputs.""" + + def get_onnx_model( + self, + inputs: dict[str, np.ndarray], + attributes: dict[str, Any], + is_constant_map: dict[str, bool], + output_dtypes: list[str], + domain_versions: dict[ONNXDomain, int], + prefix: str = "", + input_names: list[str] | None = None, + output_names: list[str] | None = None, + ) -> ModelProto: + """Build the equivalent Resize with empty optional inputs omitted.""" + del is_constant_map + if input_names is None or len(input_names) != 2: + raise ValueError("Resize rewrite requires X and sizes input names") + if output_names is None or len(output_names) != 1: + raise ValueError("Resize rewrite requires one output name") + + node_attributes = dict(attributes) + ir_version = node_attributes.pop("_ir_version", None) + resize_parameter_slot = int(node_attributes.pop("_resize_parameter_slot")) + node_inputs = [input_names[0], "", "", ""] + node_inputs[resize_parameter_slot] = input_names[1] + while node_inputs[-1] == "": + node_inputs.pop() + node = helper.make_node( + "Resize", + node_inputs, + output_names, + name=f"{prefix}Resize", + **node_attributes, + ) + + output_element_type = SupportedONNXType.from_onnx_type(output_dtypes[0]).tensor_proto_type + parameter_element_type = ( + TensorProto.INT64 if resize_parameter_slot == 3 else TensorProto.FLOAT + ) + input_specs = [ + (0, 0, output_element_type), + (resize_parameter_slot, 1, parameter_element_type), + ] + graph_inputs = [] + for node_index, schema_index, fallback_element_type in input_specs: + name = node_inputs[node_index] + value = inputs.get(_RESIZE_OMITTED_INPUTS_SCHEMA.inputs[schema_index].name) + element_type = ( + helper.np_dtype_to_tensor_dtype(value.dtype) + if value is not None + else fallback_element_type + ) + shape = list(value.shape) if value is not None else None + graph_inputs.append(helper.make_tensor_value_info(name, element_type, shape)) + + graph = helper.make_graph( + [node], + f"{prefix}ResizeWithOmittedOptionalInputs", + graph_inputs, + [ + helper.make_tensor_value_info( + output_names[0], + output_element_type, + None, + ) + ], + ) + return helper.make_model( + graph, + producer_name="winmlcli-pattern-generator", + opset_imports=[ + helper.make_opsetid(domain.schema_domain, version) + for domain, version in domain_versions.items() + ], + **({"ir_version": ir_version} if ir_version is not None else {}), + ) + + +class _ResizeAttributesPattern(_ResizeOptionalInputsPattern): + """Preserve all input slots while rewriting only selected Resize attributes.""" + + def get_schema(self) -> PatternSchema: + return _RESIZE_SCHEMA + + def check_skeleton_result( + self, + skeleton_match_result: SkeletonMatchResult, + ) -> PatternMatchResult | None: + node = skeleton_match_result.matched_nodes[0] + matcher = skeleton_match_result.matcher + if ( + matcher.domain_versions.get(ONNXDomain.AI_ONNX, 0) < 11 + or not 3 <= len(node.input) <= len(_RESIZE_SCHEMA.inputs) + or len(node.output) != 1 + ): + return None + skeleton_match_result = replace( + skeleton_match_result, + inputs=[*node.input, *([""] * (len(_RESIZE_SCHEMA.inputs) - len(node.input)))], + ) + type_mapping = self._infer_type_mapping(skeleton_match_result) + if "T1" not in type_mapping: + return None + input_types: list[int] = [] + for name in skeleton_match_result.inputs: + if not name: + input_types.append(TensorProto.UNDEFINED) + continue + tensor_type = matcher.get_tensor_type_str(name) + if not tensor_type: + return None + input_types.append(SupportedONNXType.from_onnx_type(tensor_type).tensor_proto_type) + attributes = self._infer_schema_attributes(skeleton_match_result) + attributes["_ir_version"] = matcher.model.ir_version + attributes["_input_element_types"] = input_types + return PatternMatchResult( + skeleton_match_result=skeleton_match_result, + schema_input_to_value={ + parameter.name: name + for parameter, name in zip( + _RESIZE_SCHEMA.inputs, + skeleton_match_result.inputs, + strict=True, + ) + }, + schema_output_to_value={"Y": node.output[0]}, + type_param_to_type=type_mapping, + attributes=attributes, + # No dummy arrays: only input names and types are needed for an attribute edit. + input_infos={ + parameter.name: InputInfo(name=parameter.name) + for parameter in _RESIZE_SCHEMA.inputs + }, + ) + + def _rewrite_attributes(self, attributes: dict[str, Any]) -> dict[str, Any]: + return attributes + + def get_onnx_model( + self, + inputs: dict[str, np.ndarray], + attributes: dict[str, Any], + is_constant_map: dict[str, bool], + output_dtypes: list[str], + domain_versions: dict[ONNXDomain, int], + prefix: str = "", + input_names: list[str] | None = None, + output_names: list[str] | None = None, + ) -> ModelProto: + del inputs, is_constant_map + if input_names is None or len(input_names) != len(_RESIZE_SCHEMA.inputs): + raise ValueError("Resize attribute rewrite requires all optional input slots") + if output_names is None or len(output_names) != 1: + raise ValueError("Resize attribute rewrite requires one output") + node_attributes = dict(attributes) + ir_version = node_attributes.pop("_ir_version", None) + input_types = node_attributes.pop("_input_element_types") + node_attributes = self._rewrite_attributes(node_attributes) + node_inputs = list(input_names) + while node_inputs and not node_inputs[-1]: + node_inputs.pop() + output_type = SupportedONNXType.from_onnx_type(output_dtypes[0]).tensor_proto_type + graph = helper.make_graph( + [ + helper.make_node( + "Resize", node_inputs, output_names, name=f"{prefix}Resize", **node_attributes + ) + ], + f"{prefix}ResizeAttributes", + [ + helper.make_tensor_value_info(name, dtype, None) + for name, dtype in zip(input_names, input_types, strict=True) + if name + ], + [helper.make_tensor_value_info(output_names[0], output_type, None)], + ) + return helper.make_model( + graph, + opset_imports=[ + helper.make_opsetid(domain.schema_domain, version) + for domain, version in domain_versions.items() + ], + **({"ir_version": ir_version} if ir_version is not None else {}), + ) + + +class ResizeWithTfHalfPixelForNNPattern(_ResizeAttributesPattern): + """Match nearest/floor Resize where integer scales prove coordinate equivalence.""" + + def check_skeleton_result( + self, + skeleton_match_result: SkeletonMatchResult, + ) -> PatternMatchResult | None: + """Require nearest/floor interpolation and readable non-overridable integer scales.""" + result = super().check_skeleton_result(skeleton_match_result) + if result is None: + return None + attributes = result.attributes + if ( + attributes.get("coordinate_transformation_mode") != "tf_half_pixel_for_nn" + or attributes.get("mode", "nearest") != "nearest" + or attributes.get("nearest_mode", "round_prefer_floor") != "floor" + ): + return None + matcher = skeleton_match_result.matcher + scales_name = result.schema_input_to_value["scales"] + producer = matcher.producer_lookup.get(scales_name) + if producer is None or producer[2] not in ("Initializer", "Constant"): + return None + if producer[2] == "Constant": + node = matcher.node_lookup[producer[0]] + if node.domain not in ("", "ai.onnx"): + return None + scales = matcher.tensor_values.get(scales_name) + if ( + scales is None + or scales.dtype != np.dtype(np.float32) + or scales.ndim != 1 + or scales.size == 0 + or not np.all(np.isfinite(scales) & (scales > 0)) + or not np.all(scales == np.floor(scales)) + ): + return None + sizes_name = result.schema_input_to_value["sizes"] + if sizes_name and not _is_empty_resize_input(sizes_name, matcher): + return None + input_shape = matcher.get_tensor_shape(result.schema_input_to_value["X"]) + axes = attributes.get("axes") + rank = ( + len(axes) if axes is not None else len(input_shape) if input_shape is not None else None + ) + if rank is None or scales.size != rank: + return None + return result + + +class ResizeWithAsymmetricCoordinatesPattern(_ResizeAttributesPattern): + """Keep nearest interpolation and use asymmetric coordinates.""" + + def _rewrite_attributes(self, attributes: dict[str, Any]) -> dict[str, Any]: + return {**attributes, "coordinate_transformation_mode": "asymmetric"} + + +class ResizeWithCubicInterpolationPattern(_ResizeAttributesPattern): + """Match cubic Resize without antialiasing, outside exclusion, or crop semantics.""" + + def check_skeleton_result( + self, + skeleton_match_result: SkeletonMatchResult, + ) -> PatternMatchResult | None: + """Exclude cubic settings that need separate linear-approximation handling.""" + result = super().check_skeleton_result(skeleton_match_result) + if result is None: + return None + attributes = result.attributes + if ( + attributes.get("mode") != "cubic" + or attributes.get("antialias", 0) != 0 + or attributes.get("exclude_outside", 0) != 0 + or attributes.get("coordinate_transformation_mode") == "tf_crop_and_resize" + ): + return None + return result + + +class ResizeWithLinearInterpolationPattern(_ResizeAttributesPattern): + """Approximate cubic interpolation with linear interpolation.""" + + def _rewrite_attributes(self, attributes: dict[str, Any]) -> dict[str, Any]: + return {**attributes, "mode": "linear"} + + +__all__ = [ + "ResizeWithAsymmetricCoordinatesPattern", + "ResizeWithCubicInterpolationPattern", + "ResizeWithEmptyOptionalInputsPattern", + "ResizeWithLinearInterpolationPattern", + "ResizeWithOmittedOptionalInputsPattern", + "ResizeWithTfHalfPixelForNNPattern", +] diff --git a/src/winml/modelkit/pattern/cgc/utils.py b/src/winml/modelkit/pattern/cgc/utils.py new file mode 100644 index 000000000..1e127a5f6 --- /dev/null +++ b/src/winml/modelkit/pattern/cgc/utils.py @@ -0,0 +1,52 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Static-value guards shared by explicit compatibility patterns.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from ...onnx import ONNXDomain + + +if TYPE_CHECKING: + import numpy as np + + from .. import PatternMatcher + + +def _static_tensor(name: str, matcher: PatternMatcher) -> np.ndarray | None: + """Read only direct initializers or standard Constants, never input defaults.""" + producer = matcher.producer_lookup.get(name) + if producer is None or producer[2] not in {"Initializer", "Constant"}: + return None + if producer[2] == "Constant": + node = matcher.node_lookup[producer[0]] + if node.domain not in {"", ONNXDomain.AI_ONNX.value}: + return None + return matcher.tensor_values.get(name) + + +def _depends_on_overridable_initializer(name: str, matcher: PatternMatcher) -> bool: + """Reject static-shape proofs that might depend on overridable defaults.""" + overridable = {value.name for value in matcher.graph.input} & { + value.name for value in matcher.graph.initializer + } + if not overridable: + return False + pending = [name] + visited: set[str] = set() + while pending: + current = pending.pop() + if current in visited: + continue + visited.add(current) + if current in overridable: + return True + producer = matcher.producer_lookup.get(current) + node = matcher.node_lookup.get(producer[0]) if producer else None + if node is not None: + pending.extend(node.input) + return False diff --git a/src/winml/modelkit/session/__init__.py b/src/winml/modelkit/session/__init__.py index 68017dab0..b176b6fb7 100644 --- a/src/winml/modelkit/session/__init__.py +++ b/src/winml/modelkit/session/__init__.py @@ -11,6 +11,7 @@ if TYPE_CHECKING: from ..ep_path import VALID_SOURCE_TAGS, DirectorySource, EPEntry + from ._runtime_import import import_runtime from .ep_device import ( DEVICE_TO_DEVICE_TYPE, DEVICE_TYPE_TO_DEVICE, @@ -56,11 +57,13 @@ from .monitor.qnn_monitor import QNNMonitor from .monitor.vitisai_monitor import VitisAIMonitor from .qairt.qairt_session import WinMLQairtSession + from .runtime_session import WinMLRuntimeSession from .session import InferenceError, PerfContext, SessionState, WinMLSession from .stats import PerfStats _LAZY_IMPORTS: dict[str, tuple[str, str]] = { + "import_runtime": ("._runtime_import", "import_runtime"), "VALID_SOURCE_TAGS": ("..ep_path", "VALID_SOURCE_TAGS"), "DirectorySource": ("..ep_path", "DirectorySource"), "EPEntry": ("..ep_path", "EPEntry"), @@ -111,6 +114,7 @@ "NvTensorRTRTXMonitor": (".monitor", "NvTensorRTRTXMonitor"), "VitisAIMonitor": (".monitor.vitisai_monitor", "VitisAIMonitor"), "WinMLQairtSession": (".qairt.qairt_session", "WinMLQairtSession"), + "WinMLRuntimeSession": (".runtime_session", "WinMLRuntimeSession"), "InferenceError": (".session", "InferenceError"), "PerfContext": (".session", "PerfContext"), "SessionState": (".session", "SessionState"), @@ -160,6 +164,7 @@ "WinMLEPRegistrationFailed", "WinMLEPRegistry", "WinMLQairtSession", + "WinMLRuntimeSession", "WinMLSession", "auto_detect_device", "available_eps_for_device", @@ -170,6 +175,7 @@ "ep_to_device", "eps_for_device", "expand_ep_name", + "import_runtime", "known_ep_short_names", "lookup_device_spec", "resolve_device", diff --git a/src/winml/modelkit/session/_runtime_import.py b/src/winml/modelkit/session/_runtime_import.py new file mode 100644 index 000000000..fcb09d237 --- /dev/null +++ b/src/winml/modelkit/session/_runtime_import.py @@ -0,0 +1,29 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- + +"""Lazy import of the optional Windows ML Runtime native dependency.""" + +from typing import Any + +import click + + +def import_runtime() -> Any: + """Import ``windowsml.runtime`` or raise an actionable ClickException.""" + try: + import windowsml.runtime as wr # type: ignore[import-not-found, unused-ignore] + except ImportError as e: + raise click.ClickException( + "--runtime winml-runtime requires the preview 'windowsml' package with the " + "Runtime API. Install it with the ORT backend, e.g. " + "`pip install windowsml[with-ort]`." + ) from e + except FileNotFoundError as e: # missing WinMLRuntimeCore.dll payload + raise click.ClickException( + "--runtime winml-runtime: the installed 'windowsml' package does not ship the " + "Runtime native library (WinMLRuntimeCore.dll). Install a preview build that " + "includes the Runtime API." + ) from e + return wr diff --git a/src/winml/modelkit/session/ep_device.py b/src/winml/modelkit/session/ep_device.py index 4bba1fdf4..24f1ad1a7 100644 --- a/src/winml/modelkit/session/ep_device.py +++ b/src/winml/modelkit/session/ep_device.py @@ -392,6 +392,7 @@ class EPDeviceSpec: EPDeviceSpec(ep="MIGraphXExecutionProvider", device="gpu"), EPDeviceSpec(ep="TensorrtExecutionProvider", device="gpu"), EPDeviceSpec(ep="NvTensorRTRTXExecutionProvider", device="gpu"), + EPDeviceSpec(ep="WinMLCGExecutionProvider", device="gpu"), EPDeviceSpec(ep="OpenVINOExecutionProvider", device="cpu"), # ---- QNN secondary (Snapdragon boxes without vendor-optimal alternatives) ---- EPDeviceSpec( @@ -685,7 +686,11 @@ def auto_detect_device() -> str: # --- resolution ------------------------------------------------------------ -def resolve_device(target: EPDeviceTarget) -> EPDeviceTarget: +def resolve_device( + target: EPDeviceTarget, + *, + backend: str | None = None, +) -> EPDeviceTarget: """Resolve an EP/device intent to a concrete target. Takes a typed :class:`EPDeviceTarget` intent (possibly carrying @@ -712,6 +717,9 @@ def resolve_device(target: EPDeviceTarget) -> EPDeviceTarget: target: User intent. ``target.ep`` and ``target.device`` may be the literal ``"auto"``; ``target.source`` may be ``None`` or a canonical source tag. + backend: Runtime API backend. When set, loads the Runtime native + payload before probing EP devices. The CGC backend uses DML + internally to discover device metadata. Returns: Resolved :class:`EPDeviceTarget` with no ``"auto"`` values. @@ -722,6 +730,17 @@ def resolve_device(target: EPDeviceTarget) -> EPDeviceTarget: ValueError: Unknown EP or device after deduction, or no registered EP backs the requested device. """ + if backend is not None: + from ._runtime_import import import_runtime + + import_runtime() + if backend == "cgc" and target.ep == "auto": + target = EPDeviceTarget( + ep="dml", + device=target.device, + source=None, + ) + ep = target.ep device = target.device @@ -907,6 +926,17 @@ def hardware_name(self) -> str: or "" ) + @property + def adapter_luid(self) -> int | None: + """Unsigned DXCore/DXGI adapter LUID, or ``None`` when unavailable.""" + raw = self._ort.device.metadata.get("LUID") + if not raw: + return None + try: + return int(raw, 16) if raw.lower().startswith("0x") else int(raw) + except ValueError: + raise ValueError(f"Invalid adapter LUID metadata: {raw!r}") from None + @property def vendor(self) -> str: """Hardware vendor string (e.g. ``"Intel"``) from the underlying OrtEpDevice.""" diff --git a/src/winml/modelkit/session/ep_registry.py b/src/winml/modelkit/session/ep_registry.py index 2b87bc766..b49218fc8 100644 --- a/src/winml/modelkit/session/ep_registry.py +++ b/src/winml/modelkit/session/ep_registry.py @@ -490,6 +490,11 @@ def register_ep(self, entry: EPEntry) -> WinMLEP: # to resolve — the error surfaces via the ORT exception below, # which the caller renders to the console. with _suppress_dll_load_dialogs(): + if entry.ep_name == "WinMLCGExecutionProvider": + from ._runtime_import import import_runtime + + # Preload WinMLRuntimeCore and WinMLCG's delay-loaded dependencies. + import_runtime() ort.register_execution_provider_library(arg0, str(entry.dll_path)) logger.info( "Registered EP %r from %r (arg0=%r)", entry.ep_name, entry.dll_path, arg0 diff --git a/src/winml/modelkit/session/runtime_session.py b/src/winml/modelkit/session/runtime_session.py new file mode 100644 index 000000000..63dc9ed02 --- /dev/null +++ b/src/winml/modelkit/session/runtime_session.py @@ -0,0 +1,983 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""WinML Runtime inference backend (``windowsml.runtime`` pipeline API). + +:class:`WinMLRuntimeSession` runs prebuilt ONNX and MLIR models through the +Windows ML Runtime pipeline API, exposing the same small +surface (``io_config`` / ``run`` / ``perf`` / ``device`` / ``ep_name``) as the +ORT-backed :class:`~winml.modelkit.session.session.WinMLSession`. Both satisfy +the :class:`~winml.modelkit.session.backend.InferenceBackend` protocol, so +``perf`` and ``eval`` can pick a backend uniformly through +:func:`~winml.modelkit.session.backend.create_session`. + +Unlike the ORT path (which may run the full export/optimize/quantize/compile +pipeline), this backend loads a prebuilt ``.onnx`` or ``.mlir`` directly via +``Runtime.load_model`` and drives inference through a single model stage. Output +retrieval is by tensor name for ORT-backed stages (``OrtNamedBindings.output``) +and by ordinal for non-ORT stages (``Stage.output``); both return caller-owned +NumPy copies, so ``run`` yields real named outputs for parity/eval. + +Runtime availability is optional: importing ``windowsml.runtime`` loads the +preview ``WinMLRuntimeCore.dll`` and, for ONNX models, a matching +``onnxruntime-windowsml`` ORT distribution. When either is absent the import is +guarded and surfaced as a clear, actionable :class:`click.ClickException`. +""" + +from __future__ import annotations + +import ctypes +import json +import logging +import threading +from contextlib import contextmanager +from dataclasses import dataclass +from pathlib import Path +from tempfile import TemporaryDirectory +from typing import TYPE_CHECKING, Any, cast + +import click + +from ..export.cgc.artifacts import cgc_metadata_path +from ._runtime_import import import_runtime + + +if TYPE_CHECKING: + from collections.abc import Callable, Iterator, Mapping + + import numpy as np + + from ..utils.constants import RuntimeBackend + from .ep_registry import WinMLEPDevice + from .session import PerfContext + from .stats import PerfStats + +logger = logging.getLogger(__name__) + + +# ============================================================================= +# Pure helpers (no native dependency -- unit-testable in isolation) +# ============================================================================= + +# WinML Runtime ``TensorDataType`` enum name -> NumPy dtype string. Keyed by the +# enum *name* so these helpers never import the native ``windowsml.runtime`` +# module. Only types ``Runtime.tensor_from_numpy`` round-trips are listed; +# BFLOAT16 / INT4 / UINT4 / UNDEFINED have no NumPy equivalent here. +_TENSOR_DTYPE_TO_NUMPY: dict[str, str] = { + "FLOAT32": "float32", + "FLOAT16": "float16", + "FLOAT64": "float64", + "INT8": "int8", + "UINT8": "uint8", + "INT16": "int16", + "UINT16": "uint16", + "INT32": "int32", + "UINT32": "uint32", + "INT64": "int64", + "UINT64": "uint64", + "BOOL": "bool", +} + +# ModelKit device class -> WinML Runtime ``ExecutionTargetKind`` member name. +_DEVICE_TO_KIND_NAME: dict[str, str] = {"cpu": "CPU", "gpu": "GPU", "npu": "NPU"} +_DYNAMIC_DIM_SENTINELS = {(1 << 64) - 1} + + +def _numpy_dtype_for(data_type: Any) -> str: + """Map a Runtime ``TensorDataType`` (enum or name) to a NumPy dtype string. + + Raises a clear :class:`click.ClickException` for tensor element types the + Runtime backend cannot materialize as NumPy arrays (e.g. bfloat16, int4). + """ + name = getattr(data_type, "name", str(data_type)) + dtype = _TENSOR_DTYPE_TO_NUMPY.get(name) + if dtype is None: + raise click.ClickException( + f"--runtime winml-runtime cannot benchmark a model whose I/O uses tensor " + f"data type {name!r}: it has no NumPy mapping " + "(bfloat16 / int4 / uint4 / undefined are unsupported)." + ) + return dtype + + +def _shape_with_dynamic_dims(shape: Any) -> list[int | None]: + """Normalize a Runtime schema shape to the dynamic-dim convention. + + Model schema reports a free dimension as ``0`` and post-build stage schema + reports it as ``UINT64_MAX``. Perf treats ``None`` as dynamic, so normalize + both Runtime representations. + """ + resolved: list[int | None] = [] + for dim in shape: + d = int(dim) + resolved.append(None if d <= 0 or d in _DYNAMIC_DIM_SENTINELS else d) + return resolved + + +def synth_io_config(model_schema: Any, ort_schema: Any | None) -> dict[str, Any]: + """Build an ``io_config`` dict from Runtime schema objects. + + ``model_schema`` provides ordinal ``(dtype, shape)`` descriptors + (:meth:`ModelSchema.input_desc`); ``ort_schema`` provides tensor names + (:meth:`OrtModelSchema.input_name`). When ``ort_schema`` is ``None`` (a + model that does not expose ONNX name metadata), positional + ``input_{i}`` / ``output_{i}`` names are synthesized so ordinal binding + still works. + + The returned dict matches the keys ``generate_random_inputs`` consumes + (``input_names`` / ``input_shapes`` / ``input_types``) plus output metadata + for reporting and output retrieval. + """ + input_names: list[str] = [] + input_shapes: list[list[int | None]] = [] + input_types: list[str] = [] + for i in range(model_schema.input_count): + dtype, shape = model_schema.input_desc(i) + name = ort_schema.input_name(i) if ort_schema is not None else f"input_{i}" + input_names.append(name or f"input_{i}") + input_shapes.append(_shape_with_dynamic_dims(shape)) + input_types.append(_numpy_dtype_for(dtype)) + + output_names: list[str] = [] + output_shapes: list[list[int | None]] = [] + output_types: list[str] = [] + for i in range(model_schema.output_count): + dtype, shape = model_schema.output_desc(i) + oname = ort_schema.output_name(i) if ort_schema is not None else f"output_{i}" + output_names.append(oname or f"output_{i}") + output_shapes.append(_shape_with_dynamic_dims(shape)) + output_types.append(_numpy_dtype_for(dtype)) + + return { + "input_names": input_names, + "input_shapes": input_shapes, + "input_types": input_types, + "output_names": output_names, + "output_shapes": output_shapes, + "output_types": output_types, + } + + +def _apply_io_metadata(io_config: dict[str, Any], model_path: Path) -> None: + """Preserve ONNX input ranges and restore names for CGC MLIR stages.""" + if model_path.suffix.lower() == ".onnx": + from ..onnx import get_io_config + + source_io = get_io_config(model_path) + io_config["value_ranges"] = { + name: value_range + for name, value_range in source_io["value_ranges"].items() + if name in io_config["input_names"] + } + return + + metadata_path = cgc_metadata_path(model_path) + if not metadata_path.is_file(): + return + metadata = json.loads(metadata_path.read_text(encoding="utf-8")) + input_names = [item["name"] for item in metadata.get("inputs", [])] + output_names = [item["name"] for item in metadata.get("outputs", [])] + if len(input_names) != len(io_config["input_names"]): + raise click.ClickException( + f"Runtime input count does not match {metadata_path.name}." + ) + if len(output_names) != len(io_config["output_names"]): + raise click.ClickException( + f"Runtime output count does not match {metadata_path.name}." + ) + io_config["input_names"] = input_names + io_config["output_names"] = output_names + + +def resolve_provider_kind( + device: str | None, + ep: str | None, + ep_source: str | None, + *, + resolve_fn: Callable[[str, str, str | None], tuple[str, str]] | None = None, +) -> tuple[str, str]: + """Resolve ``--device`` / ``--ep`` to a concrete ``(provider_name, device_class)``. + + By default this reuses ModelKit's :func:`resolve_device` -- the exact + resolution the ORT ``--runtime winml-ort`` path uses -- so device/EP semantics + stay consistent across backends. ``resolve_fn`` is an injection seam for + tests: it receives ``(ep, device, source)`` and returns + ``(full_ep_name, device_class)``. + """ + if resolve_fn is not None: + return resolve_fn(ep or "auto", device or "auto", ep_source) + + from .ep_device import EPDeviceTarget, resolve_device + + resolved = resolve_device( + EPDeviceTarget(ep=ep or "auto", device=device or "auto", source=ep_source) + ) + return resolved.ep, resolved.device + + +def kind_name_for_device(device_class: str) -> str: + """Map a resolved device class to an ``ExecutionTargetKind`` member name.""" + try: + return _DEVICE_TO_KIND_NAME[device_class] + except KeyError: + raise click.ClickException( + f"--runtime winml-runtime cannot target device {device_class!r}; " + f"expected one of {sorted(_DEVICE_TO_KIND_NAME)}." + ) from None + + +def _stage_schema(wr: Any, stage: Any) -> Any: + """Return the authoritative post-build stage schema. + + The Runtime projection does not yet expose ``Stage.schema()`` in all preview + wheels. Its ``IWinMLStageSchema`` ABI matches the descriptor surface wrapped + by ``ModelSchema``, so use that wrapper until the public convenience method + is available. + """ + schema_factory = getattr(stage, "schema", None) + if callable(schema_factory): + return schema_factory() + + bindings = getattr(wr, "_b", None) + schema_type = getattr(wr, "ModelSchema", None) + if bindings is None or schema_type is None: + raise click.ClickException( + "The installed windowsml Runtime projection does not expose stage schema access." + ) + + interface = stage.interface.QueryInterface(bindings.IWinMLStageSchema) + return schema_type(interface, stage) + + +def _to_numpy(value: Any) -> np.ndarray: + """Coerce a run input value (numpy array or torch tensor) to a NumPy array.""" + import numpy as np + + if isinstance(value, np.ndarray): + return value + # Duck-type torch tensors without importing torch: detach + cpu + numpy. + if hasattr(value, "detach") and hasattr(value, "cpu"): + return cast("np.ndarray", value.detach().cpu().numpy()) + if hasattr(value, "numpy"): + return cast("np.ndarray", value.numpy()) + return np.asarray(value) + + +# ============================================================================= +# Native glue (windowsml.runtime) +# ============================================================================= +class _DXCoreAdapter: + """Owned ``IDXCoreAdapter`` pointer kept alive with the Runtime target.""" + + def __init__(self, pointer: ctypes.c_void_p, module: Any) -> None: + if not pointer.value: + raise ValueError("DXCore adapter pointer must not be null.") + self._pointer = pointer + self._module = module + + @property + def pointer(self) -> int: + assert self._pointer.value is not None + return self._pointer.value + + @staticmethod + def _method( + interface: ctypes.c_void_p, + index: int, + restype: Any, + *argtypes: Any, + ) -> Any: + vtable = ctypes.cast( + interface, ctypes.POINTER(ctypes.POINTER(ctypes.c_void_p)) + ).contents + return ctypes.WINFUNCTYPE(restype, ctypes.c_void_p, *argtypes)(vtable[index]) + + @classmethod + def from_luid(cls, luid_value: int) -> _DXCoreAdapter: + """Resolve an owned adapter through ``IDXCoreAdapterFactory``.""" + from comtypes import GUID # type: ignore[import-not-found, import-untyped, unused-ignore] + + class LUID(ctypes.Structure): + _fields_ = [("LowPart", ctypes.c_uint32), ("HighPart", ctypes.c_int32)] + + if not 0 <= luid_value <= 0xFFFFFFFFFFFFFFFF: + raise click.ClickException( + f"Adapter LUID is outside uint64 range: {luid_value!r}." + ) + + try: + dxcore = ctypes.WinDLL("dxcore.dll") + except (AttributeError, OSError) as exc: + raise click.ClickException(f"Could not load dxcore.dll: {exc}") from exc + + factory_iid = GUID("{78EE5945-C36E-4B13-A669-005DD11C0F06}") + adapter_iid = GUID("{F0DB4C7F-FE5A-42A2-BD62-F2A6CF6FC83E}") + create_factory = dxcore.DXCoreCreateAdapterFactory + create_factory.argtypes = [ctypes.POINTER(GUID), ctypes.POINTER(ctypes.c_void_p)] + create_factory.restype = ctypes.c_long + + factory = ctypes.c_void_p() + hr = create_factory(ctypes.byref(factory_iid), ctypes.byref(factory)) + if hr < 0: + raise click.ClickException( + f"DXCoreCreateAdapterFactory failed (0x{hr & 0xFFFFFFFF:08X})." + ) + + adapter = ctypes.c_void_p() + try: + luid = LUID( + luid_value & 0xFFFFFFFF, + ctypes.c_int32(luid_value >> 32).value, + ) + get_adapter = cls._method( + factory, + 4, + ctypes.c_long, + ctypes.POINTER(LUID), + ctypes.POINTER(GUID), + ctypes.POINTER(ctypes.c_void_p), + ) + hr = get_adapter( + factory, + ctypes.byref(luid), + ctypes.byref(adapter_iid), + ctypes.byref(adapter), + ) + if hr < 0: + raise click.ClickException( + f"DXCore could not resolve adapter LUID {luid_value} " + f"(0x{hr & 0xFFFFFFFF:08X})." + ) + return cls(adapter, dxcore) + finally: + cls._release(factory) + + @classmethod + def _release(cls, interface: ctypes.c_void_p) -> None: + if interface.value: + cls._method(interface, 2, ctypes.c_ulong)(interface) + interface.value = None + + def close(self) -> None: + self._release(self._pointer) + self._module = None + + +@dataclass(frozen=True) +class _ResolvedRuntimeTarget: + execution_target: Any + device_class: str + provider_name: str | None = None + adapter: _DXCoreAdapter | None = None + + +# WinML native failures surface as ``windowsml.runtime.WinMLError`` carrying an +# ``hresult`` attribute. A couple of them are common enough during setup that a +# raw traceback is unhelpful, so they are translated into actionable guidance. +_HRESULT_REVISION_MISMATCH = 0x8007051A # ERROR_REVISION_MISMATCH +_HRESULT_NOT_SUPPORTED = 0x80070032 # ERROR_NOT_SUPPORTED + + +def _winml_hresult(exc: BaseException) -> int | None: + """Return a WinMLError's normalized 32-bit HRESULT, or ``None`` if not one.""" + hr = getattr(exc, "hresult", None) + if isinstance(hr, int): + return hr & 0xFFFFFFFF + return None + + +def _winml_click_error( + exc: BaseException, + hr: int, + phase: str, + *, + provider_name: str | None = None, + device_class: str | None = None, +) -> click.ClickException: + """Map a native ``(hresult, phase)`` to a clear, actionable ClickException.""" + hex_hr = f"0x{hr:08X}" + if hr == _HRESULT_REVISION_MISMATCH: + return click.ClickException( + "--runtime winml-runtime: the WinML Runtime and 'onnxruntime-windowsml' are " + f"version-incompatible ({hex_hr}). The Runtime core requires a newer ORT API " + "than the installed 'onnxruntime-windowsml' provides. Install an " + "'onnxruntime-windowsml' build that matches your 'windowsml' Runtime version." + ) + if hr == _HRESULT_NOT_SUPPORTED and phase == "build": + provider = provider_name or "the requested execution provider" + device = repr(device_class) if device_class else "the requested device" + return click.ClickException( + f"--runtime winml-runtime: building the pipeline for {provider} on {device} is " + f"not supported by this WinML Runtime build ({hex_hr}). The execution provider " + "may be unavailable in this environment (for example the ORT provider bridge " + "failed to load -- look for an 'Init provider bridge failed' warning above). " + "Try '--device cpu', or install an 'onnxruntime-windowsml' build that includes " + "this provider." + ) + label = {"load": "loading the model", "build": "building the pipeline", "run": "inference"}.get( + phase, phase + ) + return click.ClickException(f"--runtime winml-runtime: {label} failed ({hex_hr}): {exc}") + + +@contextmanager +def _translate_native_errors( + phase: str, *, provider_name: str | None = None, device_class: str | None = None +) -> Iterator[None]: + """Translate native ``WinMLError``s from a phase into actionable ClickExceptions. + + Errors we raise ourselves (``ClickException``) pass through untouched, and any + non-WinML exception (no ``hresult``) is re-raised as-is so genuine bugs stay + visible. + """ + try: + yield + except click.ClickException: + raise + except Exception as exc: + hr = _winml_hresult(exc) + if hr is None: + raise + raise _winml_click_error( + exc, hr, phase, provider_name=provider_name, device_class=device_class + ) from exc + + +def _resolve_mlir_target( + runtime: Any, + ep_device: WinMLEPDevice, +) -> _ResolvedRuntimeTarget: + """Create a Runtime hardware target from the device carried by an EP pair.""" + device = ep_device.device + device_class = device.device_type.lower() + logger.info( + "winml-runtime MLIR target resolved to %s / %s", + device_class, + device.hardware_name, + ) + if device_class == "cpu": + with _translate_native_errors("build", device_class=device_class): + target = runtime.create_cpu_target() + return _ResolvedRuntimeTarget(target, device_class) + + adapter_luid = device.adapter_luid + adapter = None + adapter_error = None + if adapter_luid is not None: + try: + adapter = _DXCoreAdapter.from_luid(adapter_luid) + except click.ClickException as exc: + adapter_error = exc + if adapter is None: + raise click.ClickException( + "Cannot resolve the selected device LUID. Run 'winml sys' and " + "select another device with --device-luid ." + ) from adapter_error + try: + with _translate_native_errors("build", device_class=device_class): + target = runtime.create_target_from_adapter(adapter.pointer) + return _ResolvedRuntimeTarget(target, device_class, adapter=adapter) + except Exception: + adapter.close() + raise + + +def _resolve_onnx_target( + runtime: Any, + wr: Any, + device: str, + ep: str | None, + ep_source: str | None, + ep_device: WinMLEPDevice | None = None, +) -> _ResolvedRuntimeTarget: + """Create an ORT-compatible Runtime target for ONNX input.""" + if ep_device is None: + provider_name, device_class = resolve_provider_kind(device, ep, ep_source) + else: + provider_name = ep_device.device.ep_name + device_class = ep_device.device.device_type.lower() + logger.info("winml-runtime target resolved to %s / %s", device_class, provider_name) + adapter: _DXCoreAdapter | None = None + with _translate_native_errors( + "build", + provider_name=provider_name, + device_class=device_class, + ): + if provider_name == "CPUExecutionProvider" and device_class == "cpu": + target = runtime.create_cpu_target() + else: + kind = getattr(wr.ExecutionTargetKind, kind_name_for_device(device_class)) + if ep_device is None or device_class == "cpu": + target = runtime.create_ort_execution_target(provider_name, kind) + else: + adapter_luid = ep_device.device.adapter_luid + if adapter_luid is None: + raise click.ClickException( + f"The selected {device_class.upper()} device " + f"{ep_device.device.hardware_name!r} does not expose " + "adapter LUID metadata." + ) + adapter = _DXCoreAdapter.from_luid(adapter_luid) + try: + hardware_target = runtime.create_target_from_adapter(adapter.pointer) + target = runtime.create_ort_execution_target( + provider_name, + kind, + hardware_target, + ) + except Exception: + adapter.close() + raise + return _ResolvedRuntimeTarget(target, device_class, provider_name, adapter) + + +def _bind_inputs( + runtime: Any, + wr: Any, + stage: Any, + input_names: list[str], + inputs: dict[str, Any], + *, + use_named_bindings: bool = True, +) -> Any: + """Bind inputs, preferring name-based ORT bindings. + + Returns the :class:`OrtNamedBindings` handle when the stage is ORT-backed + (so outputs can later be fetched by name), else ``None`` after binding by + ordinal. + """ + named = None + if use_named_bindings: + try: + named = stage.ort_bindings() + except wr.NotSupportedError: + named = None + + if named is not None: + for name in input_names: + named.bind_input(name, runtime.tensor_from_numpy(inputs[name])) + return named + + for index, name in enumerate(input_names): + stage.bind_input(index, runtime.tensor_from_numpy(inputs[name])) + return None + + +def _stage_diagnostics(wr: Any, stage: Any) -> tuple[str | None, bool]: + """Return ``(selected_provider, is_pinned)`` when the stage exposes ORT diagnostics.""" + try: + diag = stage.ort_diagnostics() + except wr.NotSupportedError: + return None, False + try: + return diag.selected_provider, diag.is_provider_pinned + except Exception: # diagnostics are advisory; never fail over them + logger.debug("Failed to read ORT stage diagnostics", exc_info=True) + return None, False + + +# ============================================================================= +# Session +# ============================================================================= +class WinMLRuntimeSession: + """WinML Runtime pipeline session for a single prebuilt model artifact. + + Mirrors the small :class:`~winml.modelkit.session.session.WinMLSession` + surface (``io_config`` / ``run`` / ``perf`` / ``device`` / ``ep_name``) but + is backed by the ``windowsml.runtime`` pipeline API instead of ORT. Native + resources (model, pipeline, stage) are built lazily on first use so + constructing a session never imports the preview native library. + """ + + def __init__( + self, + model_path: str | Path, + ep_device: WinMLEPDevice | None = None, + *, + device: str | None = "auto", + ep: str | None = None, + ep_source: str | None = None, + provider_options: Mapping[str, str] | None = None, + session_options: Callable[[], Any] | None = None, + backend: RuntimeBackend, + ) -> None: + """Initialize a Runtime session. + + Args: + model_path: Path to a prebuilt ONNX or MLIR model. + device: Device shortcut (``cpu``/``gpu``/``npu``/``auto``). + ep: Optional EP short or full name to pin (e.g. ``"qnn"``). + ep_source: Optional EP source tag (from ``--ep name@source``). + provider_options: Runtime EP options. Ignored for MLIR input and + rejected for ONNX because the Runtime pipeline API does not + expose a per-EP option hook. + backend: Resolved backend used for model execution. + """ + self._lock = threading.RLock() + self._model_path = Path(model_path) + self._is_mlir = self._model_path.suffix.lower() == ".mlir" + self._backend = backend + if self._backend == "cgc" and ep_device is None: + raise ValueError("ep_device is required for CGC Runtime sessions.") + self._ep_device = ep_device + if self._backend == "cgc": + assert ep_device is not None + device = ep_device.device.device_type.lower() + ep = None + ep_source = None + provider_options = None + elif ep_device is not None: + device = ep_device.device.device_type.lower() + ep = ep_device.ep_short_name + ep_source = ep_device.source_tag + self._device_req = device or "auto" + self._ep_req = ep + self._ep_source = ep_source + if session_options is not None: + raise ValueError( + "session_options are not supported by the Windows ML Runtime backend." + ) + if provider_options: + raise click.ClickException( + "--ep-options are not supported with --runtime winml-runtime because " + "the Windows ML Runtime API cannot apply provider options." + ) + + # Native state, built lazily by _ensure_built(). + self._wr: Any = None + self._runtime: Any = None + self._adapter_handle: _DXCoreAdapter | None = None + self._model: Any = None + self._pipeline: Any = None + self._stage: Any = None + self._io_config: dict[str, Any] | None = None + self._provider_name: str | None = None + self._device_class: str | None = None + self._selected_provider: str | None = None + self._is_pinned: bool = False + self._has_named_bindings = False + self._compiled_artifacts: TemporaryDirectory[str] | None = None + self._built = False + + # Perf tracking, enabled inside perf(). + self._perf_stats: PerfStats | None = None + + # -- lifecycle ---------------------------------------------------------- + def _ensure_built(self) -> None: + """Import the runtime, resolve the target, load and build the pipeline. + + Idempotent; native errors are translated into ClickExceptions. + """ + with self._lock: + self._ensure_built_locked() + + def _resolve_target(self, runtime: Any, wr: Any) -> _ResolvedRuntimeTarget: + if self._backend == "cgc": + assert self._ep_device is not None + return _resolve_mlir_target(runtime, self._ep_device) + return _resolve_onnx_target( + runtime, + wr, + self._device_req, + self._ep_req, + self._ep_source, + self._ep_device, + ) + + def _load_onnx_on_ort(self, runtime: Any) -> tuple[Any, Any, bool]: + """Load ONNX directly for execution by the ORT backend.""" + with _translate_native_errors("load"): + model = runtime.load_model(str(self._model_path)) + return model, model.ort_schema(), True + + def _load_onnx_on_cgc( + self, + runtime: Any, + resolved_target: _ResolvedRuntimeTarget, + ) -> tuple[Any, Any, bool]: + """Compile ONNX to CGIR and reload it for execution by the CGC backend.""" + with _translate_native_errors("load"): + source_model = runtime.load_model(str(self._model_path)) + try: + ort_schema = source_model.ort_schema() + + self._compiled_artifacts = TemporaryDirectory(prefix="winml-runtime-cgc-") + artifact_path = Path(self._compiled_artifacts.name) / "model.mlir" + with _translate_native_errors("build", device_class="gpu"): + compiler = resolved_target.execution_target.model_compiler() + try: + compiler.compile_to_file(source_model, str(artifact_path)) + finally: + compiler.close() + + with _translate_native_errors("load"): + model = runtime.load_model(str(artifact_path)) + return model, ort_schema, False + except Exception: + source_model.close() + raise + + def _load_mlir(self, runtime: Any) -> tuple[Any, None, bool]: + """Load MLIR directly for execution by the CGC backend.""" + with _translate_native_errors("load"): + return runtime.load_model(str(self._model_path)), None, False + + def _ensure_built_locked(self) -> None: + if self._built: + return + + wr = import_runtime() + runtime = wr.Runtime() + resolved_target = self._resolve_target(runtime, wr) + try: + if self._is_mlir: + model, ort_schema, has_named_bindings = self._load_mlir(runtime) + elif self._backend == "cgc": + model, ort_schema, has_named_bindings = self._load_onnx_on_cgc( + runtime, resolved_target + ) + else: + model, ort_schema, has_named_bindings = self._load_onnx_on_ort(runtime) + builder = runtime.create_pipeline_builder() + with _translate_native_errors( + "build", + provider_name=resolved_target.provider_name, + device_class=resolved_target.device_class, + ): + stage = builder.add_model_stage(model, resolved_target.execution_target) + pipeline = builder.build() + stage_schema = _stage_schema(wr, stage) + io_config = synth_io_config(stage_schema, ort_schema) + _apply_io_metadata(io_config, self._model_path) + selected_provider, is_pinned = _stage_diagnostics(wr, stage) + except Exception: + if resolved_target.adapter is not None: + resolved_target.adapter.close() + raise + + self._wr = wr + self._runtime = runtime + self._adapter_handle = resolved_target.adapter + self._model = model + self._pipeline = pipeline + self._stage = stage + self._io_config = io_config + self._provider_name = resolved_target.provider_name + self._device_class = resolved_target.device_class + self._selected_provider = selected_provider + self._is_pinned = is_pinned + self._has_named_bindings = has_named_bindings + self._built = True + + # -- metadata ----------------------------------------------------------- + @property + def io_config(self) -> dict: + """I/O metadata derived from the Runtime model schema. + + Keys mirror :attr:`WinMLSession.io_config` (``input_names`` / + ``input_shapes`` / ``input_types`` / ``output_names`` / + ``output_shapes``), plus ``output_types`` for output materialization. + """ + self._ensure_built() + assert self._io_config is not None + return self._io_config + + @property + def running_model_path(self) -> Path: + """Return the model artifact loaded by the Runtime.""" + return self._model_path + + @property + def device(self) -> str: + """Resolved Runtime target device class.""" + self._ensure_built() + assert self._device_class is not None + return self._device_class + + @property + def ep_name(self) -> str | None: + """Provider the stage resolved to, or ``None`` before the pipeline is built. + + Returns the diagnostics-reported selected provider when available, + falling back to the requested provider name (non-ORT stages expose no + diagnostics). + """ + if not self._built: + return None + return self._selected_provider or self._provider_name + + @property + def requested_provider(self) -> str | None: + """The provider name resolved from ``--device``/``--ep`` (pre-diagnostics).""" + self._ensure_built() + return self._provider_name + + @property + def is_pinned(self) -> bool: + """Whether the stage pinned the requested provider (from ORT diagnostics).""" + self._ensure_built() + return self._is_pinned + + # -- inference ---------------------------------------------------------- + def _prepare_inputs(self, inputs: dict[str, Any]) -> dict[str, np.ndarray]: + """Coerce each input to a contiguous NumPy array of the model's dtype. + + Accepts torch tensors or NumPy arrays and down/up-casts to the schema + input dtype (mirrors ``WinMLSession._prepare_inputs``), so the same feed + dict works against either backend. + """ + import numpy as np + + assert self._io_config is not None + names = self._io_config["input_names"] + dtypes = self._io_config["input_types"] + prepared: dict[str, np.ndarray] = {} + for name, want in zip(names, dtypes, strict=True): + if name not in inputs: + raise ValueError(f"Missing input {name!r}; expected inputs {list(names)}") + arr = _to_numpy(inputs[name]) + if arr.dtype != np.dtype(want): + arr = arr.astype(want) + prepared[name] = arr if arr.flags.c_contiguous else np.ascontiguousarray(arr) + return prepared + + def _read_outputs(self, named: Any) -> dict[str, np.ndarray]: + """Fetch every declared output as a caller-owned NumPy array. + + ORT-backed stages fetch by tensor name via the bound + :class:`OrtNamedBindings`; non-ORT stages fetch by ordinal. + """ + assert self._io_config is not None + output_names = self._io_config["output_names"] + outputs: dict[str, np.ndarray] = {} + if named is not None: + for name in output_names: + outputs[name] = named.output(name).to_numpy() + else: + for index, name in enumerate(output_names): + outputs[name] = self._stage.output(index).to_numpy() + return outputs + + def run(self, inputs: dict[str, Any]) -> dict[str, np.ndarray]: + """Run one inference and return ``{output_name: ndarray}``. + + Inputs are prepared (torch->numpy, dtype-cast) outside the timed region; + binding, ``Pipeline.run`` and output read-back are timed when inside a + :meth:`perf` window, keeping the measured cost comparable to the ORT + backend's ``session.run`` (feed + infer + fetch). + """ + with self._lock: + return self._run_locked(inputs) + + def _run_locked(self, inputs: dict[str, Any]) -> dict[str, np.ndarray]: + if not inputs: + raise ValueError("inputs cannot be empty") + self._ensure_built_locked() + prepared = self._prepare_inputs(inputs) + assert self._io_config is not None + input_names = self._io_config["input_names"] + output_count = len(self._io_config["output_names"]) + + def _do() -> dict[str, np.ndarray]: + named = _bind_inputs( + self._runtime, + self._wr, + self._stage, + input_names, + prepared, + use_named_bindings=self._has_named_bindings, + ) + with _translate_native_errors("run"): + for index in range(output_count): + self._stage.request_output(index) + self._pipeline.run() + return self._read_outputs(named) + + if self._perf_stats is not None: + return self._perf_stats.record(_do) + return _do() + + def compile(self) -> None: + """Build the Runtime pipeline if it has not already been built.""" + self._ensure_built() + + # -- perf --------------------------------------------------------------- + @contextmanager + def perf(self, warmup: int = 0, monitor: Any | None = None) -> Iterator[PerfContext]: + """Scoped perf window; :meth:`run` calls inside accumulate timing. + + The Runtime pipeline exposes no EP monitor / op-tracing hooks, so a + non-``None`` *monitor* is rejected with a clear error rather than + silently ignored. + """ + from .monitor.ep_monitor import NullEPMonitor + from .session import PerfContext + from .stats import PerfStats + + if monitor is not None: + raise click.ClickException( + "--monitor / op-tracing is not supported with --runtime winml-runtime; " + "the Windows ML Runtime pipeline exposes no EP monitor hooks. Re-run " + "without --monitor, or use --runtime winml-ort for op-level tracing." + ) + if self._perf_stats is not None: + raise RuntimeError( + "WinMLRuntimeSession.perf() is already active. Nested perf windows " + "are not supported." + ) + + self._ensure_built() + stats = PerfStats(warmup=warmup) + self._perf_stats = stats + try: + yield PerfContext(stats=stats, monitor=NullEPMonitor()) + finally: + self._perf_stats = None + + @property + def perf_stats(self) -> Any: + """Active :class:`PerfStats` inside a :meth:`perf` window, else ``None``.""" + return self._perf_stats + + # -- teardown ----------------------------------------------------------- + def close(self) -> None: + """Release native handles (best-effort; safe to call more than once).""" + with self._lock: + self._close_locked() + + def _close_locked(self) -> None: + for attr in ("_pipeline", "_stage", "_model", "_runtime", "_adapter_handle"): + obj = getattr(self, attr, None) + if obj is None: + continue + closer = getattr(obj, "close", None) + if callable(closer): + try: + closer() + except Exception: # teardown is best-effort + logger.debug("Failed to close runtime %s", attr, exc_info=True) + setattr(self, attr, None) + if self._compiled_artifacts is not None: + self._compiled_artifacts.cleanup() + self._compiled_artifacts = None + self._built = False + + def reset(self) -> None: + """Release the pipeline so a later call can rebuild it.""" + with self._lock: + self._close_locked() + self._wr = None + self._io_config = None + self._provider_name = None + self._device_class = None + self._selected_provider = None + self._is_pinned = False + self._has_named_bindings = False + + def __del__(self) -> None: + try: + self.close() + except Exception: + # Finalization must tolerate partial construction and interpreter shutdown. + pass diff --git a/src/winml/modelkit/utils/cli.py b/src/winml/modelkit/utils/cli.py index a69671273..eb71f9b2d 100644 --- a/src/winml/modelkit/utils/cli.py +++ b/src/winml/modelkit/utils/cli.py @@ -9,8 +9,9 @@ import json import os import re +from dataclasses import fields, is_dataclass from pathlib import Path -from typing import TYPE_CHECKING, Any, Literal, TypeAlias, TypedDict, TypeVar +from typing import TYPE_CHECKING, Any, Literal, TypeAlias, TypedDict, TypeVar, get_type_hints import click from rich.console import Console @@ -32,6 +33,9 @@ # Allowed values for ``--format`` / ``-f``. OutputFormat: TypeAlias = Literal["text", "json", "table", "compact"] +PrimitiveOptionType: TypeAlias = type[bool] | type[int] | type[float] | type[str] +PrimitiveOptionValue: TypeAlias = bool | int | float | str +OptionsT = TypeVar("OptionsT") class _CacheExtraKwargs(TypedDict): @@ -536,6 +540,74 @@ def parse_ep_options(values: tuple[str, ...]) -> dict[str, str] | None: return options +def parse_options( + values: tuple[str, ...], + options_type: type[OptionsT], + *, + param_hint: str = "--options", +) -> OptionsT: + """Parse repeatable ``KEY=VALUE`` options into a typed dataclass. + + Field names use CLI spelling with hyphens in place of underscores. Supported + field types are ``bool``, ``int``, ``float``, and ``str``. Unspecified + fields retain their dataclass defaults. + """ + if not is_dataclass(options_type): + raise TypeError("options_type must be a dataclass type") + + type_hints = get_type_hints(options_type) + schema: dict[str, PrimitiveOptionType] = {} + for field in fields(options_type): + value_type = type_hints[field.name] + if value_type not in {bool, int, float, str}: + raise TypeError( + f"Unsupported option type for '{field.name}': {value_type!r}" + ) + schema[field.name] = value_type + + options: dict[str, PrimitiveOptionValue] = {} + for item in values: + if "=" not in item: + raise click.BadParameter( + f"Invalid option format: '{item}'. Use KEY=VALUE.", + param_hint=param_hint, + ) + + cli_key, raw_value = (part.strip() for part in item.split("=", 1)) + if not cli_key: + raise click.BadParameter( + f"Invalid option format: '{item}'. Key cannot be empty.", + param_hint=param_hint, + ) + key = cli_key.replace("-", "_") + if key not in schema: + raise click.BadParameter( + f"Unsupported option: '{cli_key}'.", + param_hint=param_hint, + ) + + value_type = schema[key] + if value_type is bool: + normalized = raw_value.lower() + if normalized not in {"true", "false"}: + raise click.BadParameter( + f"Option '{cli_key}' expects true or false, got '{raw_value}'.", + param_hint=param_hint, + ) + options[key] = normalized == "true" + elif value_type in {int, float, str}: + try: + options[key] = value_type(raw_value) + except ValueError as e: + raise click.BadParameter( + f"Option '{cli_key}' expects {value_type.__name__}, " + f"got '{raw_value}'.", + param_hint=param_hint, + ) from e + + return options_type(**options) + + def device_option( required: bool = True, optional_message: str | None = None, diff --git a/src/winml/modelkit/utils/constants.py b/src/winml/modelkit/utils/constants.py index b536c31dd..f15389eea 100644 --- a/src/winml/modelkit/utils/constants.py +++ b/src/winml/modelkit/utils/constants.py @@ -6,6 +6,7 @@ from __future__ import annotations +from pathlib import Path from typing import Any, Literal, TypeAlias, cast, get_args, overload @@ -19,7 +20,9 @@ "EP_ALIAS_NAMES", "EP_NAMES", "EP_SUPPORTED_DEVICES", + "EXPORT_TARGETS", "ORT_SESSION_COMPILER", + "RUNTIME_BACKENDS", "RUNTIME_NAMES", "SUPPORTED_DEVICES", "SUPPORTED_EPS", @@ -28,9 +31,12 @@ "EPAlias", "EPName", "EPNameOrAlias", + "ExportTarget", + "RuntimeBackend", "RuntimeName", "extract_ep_options", "normalize_ep_name", + "resolve_runtime_api_backend", ] @@ -46,6 +52,7 @@ "QNNExecutionProvider", "TensorrtExecutionProvider", "VitisAIExecutionProvider", + "WinMLCGExecutionProvider", ] # Shorthand aliases users can pass on the CLI (case-insensitive at the parser layer). @@ -60,6 +67,7 @@ "nv_tensorrt_rtx", "migraphx", "tensorrt", + "winmlcg", ] # Either an alias or a full name — what user-facing entry points accept before normalization. @@ -80,12 +88,42 @@ # Runtime-iterable form of ``CompilerName`` (e.g. for the CLI choice list). COMPILER_NAMES: tuple[CompilerName, ...] = get_args(CompilerName) - # Inference runtimes selectable via ``winml perf --runtime``. -RuntimeName = Literal["auto", "winml-ort", "ort-genai"] +RuntimeName = Literal["auto", "winml-ort", "ort-genai", "winml-runtime"] RUNTIME_NAMES: tuple[RuntimeName, ...] = get_args(RuntimeName) +RuntimeBackend = Literal["ort", "cgc"] +RUNTIME_BACKENDS: tuple[RuntimeBackend, ...] = get_args(RuntimeBackend) + + +def resolve_runtime_api_backend( + runtime: str, + model_path: object, + backend: RuntimeBackend | None = None, +) -> RuntimeBackend | None: + """Validate and resolve the backend used by the Windows ML Runtime API.""" + if backend is not None and backend not in RUNTIME_BACKENDS: + raise ValueError( + f"Invalid Runtime API backend {backend!r}; expected one of {RUNTIME_BACKENDS}." + ) + if runtime != "winml-runtime": + if backend is not None: + raise ValueError("--backend is only supported with --runtime winml-runtime.") + return None + is_mlir = ( + isinstance(model_path, (str, Path)) + and Path(model_path).suffix.lower() == ".mlir" + ) + if is_mlir and backend == "ort": + raise ValueError("MLIR inputs require the CGC backend.") + return "cgc" if is_mlir or backend is None else backend + +# Output formats selectable via ``winml export --target``. +ExportTarget = Literal["onnx", "cgir"] +EXPORT_TARGETS: tuple[ExportTarget, ...] = get_args(ExportTarget) + + # Supported execution providers — derived from the ``EPName`` Literal above so # that ``utils.constants`` stays leaf-level (no import dependency on sysinfo). # Membership parity with ``sysinfo.device._EP_DEVICE_MAP`` is enforced by @@ -104,6 +142,7 @@ "nv_tensorrt_rtx": "NvTensorRTRTXExecutionProvider", "migraphx": "MIGraphXExecutionProvider", "tensorrt": "TensorrtExecutionProvider", + "winmlcg": "WinMLCGExecutionProvider", } # Runtime-iterable forms of the Literal types above (for membership checks, choice lists). @@ -219,6 +258,7 @@ def extract_ep_options(kwargs: dict) -> dict[str, str]: "OpenVINOExecutionProvider": ("npu", "gpu", "cpu"), "TensorrtExecutionProvider": ("gpu",), "DmlExecutionProvider": ("gpu",), + "WinMLCGExecutionProvider": ("gpu",), "CPUExecutionProvider": ("cpu",), "VitisAIExecutionProvider": ("npu",), } diff --git a/tests/unit/analyze/test_static_analyzer_cli.py b/tests/unit/analyze/test_static_analyzer_cli.py index 2a06045e6..aabf15d0f 100644 --- a/tests/unit/analyze/test_static_analyzer_cli.py +++ b/tests/unit/analyze/test_static_analyzer_cli.py @@ -1771,6 +1771,7 @@ def test_auto_ep_concrete_and_all_device_share_same_exact_local_ranking( ("OpenVINOExecutionProvider", "CPU"), ("TensorrtExecutionProvider", "GPU"), ("DmlExecutionProvider", "GPU"), + ("WinMLCGExecutionProvider", "GPU"), ("CPUExecutionProvider", "CPU"), ("VitisAIExecutionProvider", "NPU"), ], diff --git a/tests/unit/build/test_policy.py b/tests/unit/build/test_policy.py new file mode 100644 index 000000000..643b1a385 --- /dev/null +++ b/tests/unit/build/test_policy.py @@ -0,0 +1,86 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Tests for target-driven CGC build configuration.""" + +import pytest + +from winml.modelkit.config import WinMLBuildConfig +from winml.modelkit.config.build import _apply_cgc_config +from winml.modelkit.export import WinMLExportConfig +from winml.modelkit.optim import WinMLOptimizationConfig +from winml.modelkit.utils.constants import RuntimeBackend + + +@pytest.mark.parametrize( + ("backend", "ep"), + [ + (None, "cpu"), + ("ort", "openvino"), + ("ort", "nvtensorrtrtx"), + ], +) +def test_non_cgc_target_preserves_build_stages( + backend: RuntimeBackend | None, ep: str +) -> None: + config = WinMLBuildConfig() + original_export = config.export + original_optim = config.optim + original_quant = config.quant + original_compile = config.compile + original_auto = config.auto + + _apply_cgc_config(config, backend=backend, ep=ep) + + assert config.export is original_export + assert config.optim is original_optim + assert config.quant is original_quant + assert config.compile is original_compile + assert config.auto is original_auto + assert config.skip_optimize is False + assert config.convert is None + + +@pytest.mark.parametrize( + ("backend", "ep"), + [ + ("cgc", None), + ("cgc", "winmlcg"), + ("cgc", "WinMLCGExecutionProvider"), + (None, "winmlcg"), + (None, "WinMLCGExecutionProvider"), + ], +) +def test_cgc_target_enables_compatibility_optimization( + backend: RuntimeBackend | None, ep: str | None +) -> None: + config = WinMLBuildConfig() + original_export = config.export + original_quant = config.quant + + _apply_cgc_config(config, backend=backend, ep=ep) + + assert config.export is original_export + assert config.quant is original_quant + assert config.auto is False + assert config.skip_optimize is False + assert config.optim == WinMLOptimizationConfig.for_cgc() + assert config.compile is None + + if backend == "cgc": + assert config.convert is not None + assert config.convert.target == "cgir" + assert not config.convert.options + else: + assert config.convert is None + + +def test_cgc_target_preserves_existing_conversion_config() -> None: + conversion = WinMLExportConfig(target="cgir", options={}) + config = WinMLBuildConfig(convert=conversion) + + _apply_cgc_config(config, backend="cgc", ep=None) + + assert config.convert is conversion + assert config.convert.options == {} diff --git a/tests/unit/commands/test_build.py b/tests/unit/commands/test_build.py index 37f6bde79..753e2c7b3 100644 --- a/tests/unit/commands/test_build.py +++ b/tests/unit/commands/test_build.py @@ -21,6 +21,54 @@ from winml.modelkit.session import EPDeviceTarget +@pytest.mark.parametrize("is_onnx", [False, True]) +@pytest.mark.parametrize("reused", [False, True]) +@pytest.mark.parametrize("convert_enabled", [False, True]) +def test_cli_convert_stage(tmp_path, is_onnx, reused, convert_enabled): + from winml.modelkit.commands.build import _run_single_build + from winml.modelkit.config import WinMLBuildConfig + from winml.modelkit.export import WinMLExportConfig + from winml.modelkit.export.cgc import CGCExporter, CGCOptions + + config = WinMLBuildConfig() + if convert_enabled: + config.convert = WinMLExportConfig(target="cgir", options={"external_weights": True}) + config = WinMLBuildConfig.from_dict(config.to_dict()) + compile_config = config.compile + timings = None if reused else [("Optimize", 0.0)] + with ( + patch("winml.modelkit.commands.build._build_hf_pipeline", return_value=timings) as hf, + patch("winml.modelkit.commands.build._build_onnx_pipeline", return_value=timings) as onnx, + patch.object(CGCExporter, "export_onnx", autospec=True) as convert, + ): + _run_single_build( + config=config, + config_file=None, + model_id=str(tmp_path / "source.onnx") if is_onnx else "test-model", + is_onnx=is_onnx, + resolved_dir=tmp_path, + rebuild=False, + cache_key=None, + ep=None, + device=None, + extra_kwargs={}, + ) + selected, unused = (onnx, hf) if is_onnx else (hf, onnx) + selected.assert_called_once() + unused.assert_not_called() + assert config.compile is compile_config + if convert_enabled: + convert.assert_called_once() + assert convert.call_args.kwargs == { + "model": tmp_path / "model.onnx", + "output_path": tmp_path / "model.mlir", + } + assert convert.call_args.args[0].options == CGCOptions(**config.convert.options) + else: + convert.assert_not_called() + assert "convert" not in config.to_dict() + + _DEVICE_TO_EPS = { "npu": ["QNNExecutionProvider"], "gpu": ["DmlExecutionProvider"], @@ -2343,6 +2391,7 @@ def test_pre_quantized_stamp_runs_before_optimize(self, tmp_path: Path) -> None: output_dir = tmp_path / "out" config = MagicMock() + config.is_cgc = False config.skip_optimize = False config.quant = MagicMock(name="quant_config") config.validate.return_value = None diff --git a/tests/unit/commands/test_cgc_target_args.py b/tests/unit/commands/test_cgc_target_args.py new file mode 100644 index 000000000..138251ce8 --- /dev/null +++ b/tests/unit/commands/test_cgc_target_args.py @@ -0,0 +1,20 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""CGC backend and EP selection are mutually exclusive at the CLI boundary.""" + +import pytest +from click.testing import CliRunner + +from winml.modelkit.commands.build import build +from winml.modelkit.commands.config import config + + +@pytest.mark.parametrize("command", [config, build], ids=["config", "build"]) +@pytest.mark.parametrize("ep", ["dml", "winmlcg", "DmlExecutionProvider", "dml@bundled"]) +def test_cgc_backend_rejects_ep_before_model_resolution(command, ep): + result = CliRunner().invoke(command, ["--backend", "cgc", "--ep", ep]) + + assert result.exit_code == 2 + assert "--backend cgc cannot be combined with --ep." in result.output diff --git a/tests/unit/commands/test_config_cli.py b/tests/unit/commands/test_config_cli.py index 0c390f6ce..5f9ccf08d 100644 --- a/tests/unit/commands/test_config_cli.py +++ b/tests/unit/commands/test_config_cli.py @@ -118,6 +118,98 @@ def mock_generate_config(): # ============================================================================= +@pytest.mark.parametrize("backend", [None, "ort", "cgc"]) +@pytest.mark.parametrize("no_quant", [False, True]) +@pytest.mark.parametrize("kind", ["hf", "onnx", "module"]) +def test_backend_stage_settings( + runner, tmp_path, onnx_model_path, mock_generate_config, backend, no_quant, kind +): + from winml.modelkit.commands.config import config + from winml.modelkit.config import WinMLBuildConfig, resolve_quant_compile_config + from winml.modelkit.config.build import _apply_cgc_config + + cfg = WinMLBuildConfig.from_dict(mock_generate_config.return_value.to_dict()) + if kind == "onnx": + cfg.export = None + original = cfg.to_dict() + cfg.quant, cfg.compile = resolve_quant_compile_config(device="gpu", backend=backend) + _apply_cgc_config(cfg, backend=backend, ep=None) + mock_generate_config.return_value = [cfg] if kind == "module" else cfg + output = tmp_path / "config.json" + args = [ + "-m", str(onnx_model_path) if kind == "onnx" else "test-model", + "-o", str(output), + ] + if kind == "module": + args += ["--module", "test-module"] + if backend is not None: + args += ["--backend", backend] + if no_quant: + args += ["--no-quant"] + with patch("winml.modelkit.config.generate_onnx_build_config", return_value=cfg) as generate: + result = runner.invoke(config, args) + assert result.exit_code == 0, result.output + generator = generate if kind == "onnx" else mock_generate_config + assert generator.call_args.kwargs["backend"] == backend + data = json.loads(output.read_text(encoding="utf-8")) + if kind == "module": + data = data[0] + assert data.get("loader") == original.get("loader") + assert data["export"] == original["export"] + if backend == "cgc": + assert data["auto"] is False + _assert_cgc_optim(data["optim"]) + assert data["compile"] is None + assert data["convert"]["target"] == "cgir" + if no_quant: + assert data["quant"] is None + else: + assert data["quant"]["mode"] == "fp16" + else: + assert data == original + + +@pytest.mark.parametrize( + ("backend", "ep"), + [(None, None), ("ort", None), ("cgc", None), (None, "winmlcg"), + (None, "WinMLCGExecutionProvider")], +) +def test_generator_cgc_settings(onnx_model_path, backend, ep): + from winml.modelkit.config import generate_build_config + + cfg = generate_build_config( + onnx_path=onnx_model_path, device="gpu", backend=backend, ep=ep + ) + assert cfg.export is None + if backend == "cgc" or ep is not None: + assert cfg.auto is False + _assert_cgc_optim(cfg.optim) + assert cfg.quant.mode == "fp16" + assert cfg.compile is None + else: + assert cfg.auto is True + assert cfg.quant is None + if backend == "cgc": + assert cfg.convert.target == "cgir" + else: + assert cfg.convert is None + + +def _assert_cgc_optim(optim): + from winml.modelkit.optim.pipes import CGIRRewritePipe, ORTGraphPipe + + assert optim["ort_graph_optimization"] is False + assert "backend" not in optim + rules = CGIRRewritePipe.build_config(**optim).rules + all_options = { + capability.python_name: True + for capability in CGIRRewritePipe.capabilities.values() + } + assert rules == CGIRRewritePipe.build_config(**all_options).rules + assert len(optim) - 1 == len(rules) + assert not ORTGraphPipe.should_process(ORTGraphPipe.build_config(**optim)) + + class TestConfigCliInterface: """Test CLI flag parsing and help text.""" diff --git a/tests/unit/commands/test_eval.py b/tests/unit/commands/test_eval.py index d9141b7e1..7a0923632 100644 --- a/tests/unit/commands/test_eval.py +++ b/tests/unit/commands/test_eval.py @@ -98,6 +98,18 @@ def test_plain_onnx_with_model_id(self, onnx_file): assert path == str(onnx_file) assert mid == "microsoft/resnet-50" + def test_plain_mlir_with_model_id(self, tmp_path): + mlir_file = tmp_path / "model.mlir" + mlir_file.write_text("module {}") + + path, mid = _resolve_model_path( + model=(str(mlir_file),), + model_id="microsoft/resnet-50", + ) + + assert path == str(mlir_file) + assert mid == "microsoft/resnet-50" + def test_plain_onnx_without_model_id_raises(self, onnx_file): with pytest.raises(click.UsageError, match="--model-id is required"): _resolve_model_path(model=(str(onnx_file),), model_id=None) @@ -427,7 +439,8 @@ def test_model_help_mentions_onnx_model_id_and_role_path(self, runner: CliRunner result = runner.invoke(eval_cmd, ["--help"]) assert result.exit_code == 0, result.output - assert "requires --model-id" in result.output + assert "--model-id" in result.output + assert "MLIR" in result.output assert "role=path" in result.output def test_help_mentions_input_data(self, runner: CliRunner): @@ -481,6 +494,18 @@ def test_help_mentions_cache_controls(self, runner: CliRunner): assert "--use-cache / --no-use-cache" in result.output assert "--rebuild / --no-rebuild" in result.output + def test_backend_rejected_for_other_runtime(self, runner: CliRunner): + from winml.modelkit.commands.eval import eval as eval_cmd + + result = runner.invoke( + eval_cmd, + ["-m", "test/model", "--runtime", "winml-ort", "--backend", "cgc"], + obj={"debug": False}, + ) + + assert result.exit_code == 2 + assert "--backend is only supported with --runtime winml-runtime" in result.output + class TestResolveReference: def test_none_is_noop(self): @@ -497,13 +522,13 @@ def test_happy_path(self, onnx_file, onnx_vision): _resolve_reference(cfg) assert cfg.reference_path == str(onnx_vision) - def test_requires_onnx_candidate(self): + def test_requires_model_file_candidate(self): cfg = WinMLEvaluationConfig( model_path=None, reference_path="ref.onnx", mode="compare", ) - with pytest.raises(click.UsageError, match="single ONNX file"): + with pytest.raises(click.UsageError, match="single model file"): _resolve_reference(cfg) def test_composite_candidate_rejected(self, onnx_vision): @@ -512,7 +537,7 @@ def test_composite_candidate_rejected(self, onnx_vision): reference_path=str(onnx_vision), mode="compare", ) - with pytest.raises(click.UsageError, match="single ONNX file"): + with pytest.raises(click.UsageError, match="single model file"): _resolve_reference(cfg) def test_non_onnx_suffix_raises(self, onnx_file, tmp_path): @@ -1722,10 +1747,26 @@ def test_composite_model_path_dict_renders_readable_not_dict_repr(self): # Header joins the sub-model paths; detail lines list them per role ... assert "enc.onnx" in text assert "dec.onnx" in text - assert "ONNX (encoder):" in text + assert "Model (encoder):" in text # ... and never leak a raw Python dict repr. assert "{'encoder'" not in text + def test_compare_runtime_ort_shows_candidate_ep(self): + from winml.modelkit.eval import WinMLEvaluationConfig + + text = self._render( + WinMLEvaluationConfig( + model_path="candidate.onnx", + reference_path="reference.onnx", + runtime="winml-runtime", + backend="ort", + ep="dml", + mode="compare", + ) + ) + + assert "Candidate EP: dml" in text + def test_two_onnx_compare_shows_candidate_path_without_model_id(self): from winml.modelkit.eval import WinMLEvaluationConfig @@ -1737,7 +1778,14 @@ def test_two_onnx_compare_shows_candidate_path_without_model_id(self): ) ) assert "Evaluation: cand.onnx" in text - assert "ref.onnx" in text + assert "Candidate: cand.onnx" in text + assert "Candidate runtime: winml-ort" in text + assert "Candidate device: auto" in text + assert "Candidate EP: auto" in text + assert "Reference: ref.onnx" in text + assert "Reference runtime: winml-ort" in text + assert "Reference device: cpu" in text + assert "Reference EP: auto" in text # --------------------------------------------------------------------------- diff --git a/tests/unit/commands/test_eval_pytorch.py b/tests/unit/commands/test_eval_pytorch.py index 86d959d82..3c65cc3af 100644 --- a/tests/unit/commands/test_eval_pytorch.py +++ b/tests/unit/commands/test_eval_pytorch.py @@ -24,7 +24,7 @@ def test_help_shows_runtime_choices(self) -> None: result = CliRunner().invoke(eval, ["--help"]) assert result.exit_code == 0 - assert "--runtime [winml-ort|pytorch]" in result.output + assert "--runtime [winml-ort|winml-runtime|pytorch]" in result.output def test_pytorch_runtime_dispatches_pytorch(self, tmp_path) -> None: captured: dict[str, WinMLEvaluationConfig] = {} diff --git a/tests/unit/commands/test_export.py b/tests/unit/commands/test_export.py index 0c2918895..a9eda7b73 100644 --- a/tests/unit/commands/test_export.py +++ b/tests/unit/commands/test_export.py @@ -77,6 +77,179 @@ def test_export_help_shows_all_options(self, runner: CliRunner) -> None: assert "--export-config" in result.output assert "--dynamic-axes" in result.output assert "--submodel" in result.output + assert "--target" in result.output + assert "--options" in result.output + assert "--additional-options" not in result.output + assert "--exporter" not in result.output + assert "--cgc-option" not in result.output + + def test_cgc_guards_external_weight_sidecar_before_export( + self, + runner: CliRunner, + tmp_path: Path, + ) -> None: + from winml.modelkit.commands.export import export + from winml.modelkit.export.cgc import CGCExporter + + source = tmp_path / "source.onnx" + output = tmp_path / "model.mlir" + output.with_name(f"{output.name}.data").write_bytes(b"existing") + source.write_bytes(b"onnx") + + with patch.object(CGCExporter, "export_onnx") as backend_export: + result = runner.invoke( + export, + [ + "--model", + str(source), + "--output", + str(output), + "--target", + "cgir", + "--options", + "external-weights=true", + ], + ) + + assert result.exit_code != 0 + assert "Output sidecar" in result.output + assert "model.mlir.data" in result.output + backend_export.assert_not_called() + + @pytest.mark.parametrize("onnx_input", [False, True]) + @pytest.mark.parametrize("config_flag", ["cli", "--export-config", "-c", "both"]) + @pytest.mark.parametrize("override", [False, True]) + def test_target_config_and_cli_option_precedence( + self, runner, tmp_path, onnx_input, config_flag, override + ): + from winml.modelkit.commands.export import export + from winml.modelkit.export import WinMLExportConfig + from winml.modelkit.export.cgc import CGCExporter, CGCOptions + + settings = { + "target": "cgir", + "options": {"external_weights": "invalid" if override else True, "update_opset": False}, + "dynamo": True, + } + config_path = tmp_path / "export.json" + config_path.write_text( + json.dumps({"export": settings} if config_flag == "-c" else settings) + ) + source = tmp_path / "source.onnx" + source.write_bytes(b"onnx") + args = [ + "-m", + str(source) if onnx_input else "test-model", + "-o", + str(tmp_path / "model.mlir"), + ] + if config_flag == "cli": + args += ["--target", settings["target"]] + if not override: + for key, value in settings["options"].items(): + args += ["--options", f"{key}={value}"] + elif config_flag == "both": + build_path = tmp_path / "build.json" + build_path.write_text(json.dumps({"export": settings})) + config_path.write_text(json.dumps({"options": {"topo_sort_nodes": False}})) + args += ["-c", str(build_path), "--export-config", str(config_path)] + else: + args += [config_flag, str(config_path)] + if override: + args += ["--options", "external-weights=true", "--options", "external-weights=false"] + method = "export_onnx" if onnx_input else "export_pytorch" + with ( + patch.object(CGCExporter, method, autospec=True) as backend, + patch("winml.modelkit.export.export_pytorch") as onnx_backend, + patch("winml.modelkit.loader.load_hf_model", return_value=(MagicMock(), None, None)), + patch( + "winml.modelkit.loader.resolution.resolve_composite_components", + return_value=None, + ), + patch( + "winml.modelkit.export.resolve_export_config", + return_value=(WinMLExportConfig(), None), + ), + ): + result = runner.invoke(export, args) + assert result.exit_code == 0, result.output + backend.assert_called_once() + onnx_backend.assert_not_called() + assert backend.call_args.args[0].options == CGCOptions( + external_weights=not override and config_flag != "both", + update_opset=override or config_flag == "both", + topo_sort_nodes=override or config_flag != "both", + ) + if not onnx_input: + config = backend.call_args.kwargs["export_config"] + assert config.target == ("onnx" if config_flag == "cli" else settings["target"]) + assert config.dynamo is (False if config_flag == "cli" else settings["dynamo"]) + + @pytest.mark.parametrize("target", ["onnx", "cgir"]) + @pytest.mark.parametrize("config_flag", ["-c", "--export-config"]) + def test_explicit_target_overrides_export_config(self, runner, tmp_path, target, config_flag): + from winml.modelkit.commands.export import export + from winml.modelkit.export import WinMLExportConfig + from winml.modelkit.export.cgc import CGCExporter + from winml.modelkit.utils.constants import EXPORT_TARGETS + + config_path = tmp_path / "export.json" + configured_target = next(t for t in EXPORT_TARGETS if t != target) + settings = {"target": configured_target} + config_path.write_text( + json.dumps({"export": settings} if config_flag == "-c" else settings) + ) + with ( + patch.object(CGCExporter, "export_pytorch") as cgc_backend, + patch("winml.modelkit.export.export_pytorch") as onnx_backend, + patch("winml.modelkit.loader.load_hf_model", return_value=(MagicMock(), None, None)), + patch( + "winml.modelkit.loader.resolution.resolve_composite_components", + return_value=None, + ), + patch( + "winml.modelkit.export.resolve_export_config", + return_value=(WinMLExportConfig(), None), + ), + ): + result = runner.invoke(export, [ + "-m", "test-model", "-o", str(tmp_path / "model.out"), + config_flag, str(config_path), "--target", target, + ]) + assert result.exit_code == 0, result.output + selected, unused = ( + (onnx_backend, cgc_backend) if target == "onnx" else (cgc_backend, onnx_backend) + ) + selected.assert_called_once() + unused.assert_not_called() + assert selected.call_args.kwargs["export_config"].target == configured_target + + @pytest.mark.parametrize( + "settings", + [ + {"target": "invalid"}, + {"options": {"external_weights": True}}, + {"target": "cgir", "options": []}, + {"target": "cgir", "options": {"unknown": True}}, + {"target": "cgir", "options": {"external_weights": "invalid"}}, + ], + ) + def test_invalid_target_config_rejected_before_onnx_export(self, runner, tmp_path, settings): + from winml.modelkit.commands.export import export + from winml.modelkit.export.cgc import CGCExporter + + config_path = tmp_path / "export.json" + config_path.write_text(json.dumps(settings)) + source = tmp_path / "source.onnx" + source.write_bytes(b"onnx") + with patch.object(CGCExporter, "export_onnx") as backend: + result = runner.invoke(export, [ + "-m", str(source), "-o", str(tmp_path / "model.mlir"), + "--export-config", str(config_path), + ]) + assert result.exit_code != 0 + assert "requires a HuggingFace model ID" not in result.output + backend.assert_not_called() def test_export_help_examples_run(self, runner: CliRunner, tmp_path: Path) -> None: """Every command example in export help should execute without crashing.""" @@ -1267,6 +1440,144 @@ def fake_export_onnx(**kwargs): assert "1 sub-model" in result.output +class TestExportCGC: + """Test CGC uses the shared export orchestration.""" + + def test_existing_onnx_uses_direct_cgc_entrypoint( + self, + runner: CliRunner, + tmp_path: Path, + ) -> None: + from winml.modelkit.commands.export import export + from winml.modelkit.export.cgc import CGCExporter + + source = tmp_path / "source.onnx" + output = tmp_path / "model.mlir" + source.write_bytes(b"onnx") + + with ( + patch.object(CGCExporter, "export_onnx") as export_onnx, + patch.object(CGCExporter, "export_pytorch") as export_pytorch, + ): + result = runner.invoke( + export, + [ + "--model", + str(source), + "--output", + str(output), + "--target", + "cgir", + ], + ) + + assert result.exit_code == 0, result.output + export_onnx.assert_called_once_with(model=source, output_path=output) + export_pytorch.assert_not_called() + + def test_composite_uses_cgc_entrypoint_per_component( + self, + runner: CliRunner, + mock_export_onnx: MagicMock, + tmp_path: Path, + ) -> None: + from winml.modelkit.commands.export import export + from winml.modelkit.export import WinMLExportConfig + from winml.modelkit.export.cgc import CGCExporter + from winml.modelkit.loader import WinMLLoaderConfig + + components = { + "image-encoder": "image-feature-extraction", + "text-encoder": "feature-extraction", + } + output_path = tmp_path / "clip.mlir" + + with ( + patch( + "winml.modelkit.loader.resolution.resolve_composite_components", + return_value=components, + ), + patch( + "winml.modelkit.loader.load_hf_model", + side_effect=lambda _model, task=None, **_kwargs: (MagicMock(), None, task), + ), + patch( + "winml.modelkit.export.resolve_export_config", + return_value=( + WinMLExportConfig(), + WinMLLoaderConfig(task="zero-shot-image-classification"), + ), + ), + patch.object(CGCExporter, "export_pytorch") as export_pytorch, + ): + result = runner.invoke( + export, + [ + "--model", + "openai/clip-vit-base-patch32", + "--task", + "zero-shot-image-classification", + "--output", + str(output_path), + "--target", + "cgir", + ], + obj={"debug": False}, + ) + + assert result.exit_code == 0, result.output + assert export_pytorch.call_count == len(components) + assert {Path(call.kwargs["output_path"]) for call in export_pytorch.call_args_list} == { + output_path.with_stem(f"{output_path.stem}_{name}") for name in components + } + assert {call.kwargs["task"] for call in export_pytorch.call_args_list} == set( + components.values() + ) + mock_export_onnx.assert_not_called() + + def test_composite_guards_all_cgc_sidecars_before_export( + self, + runner: CliRunner, + tmp_path: Path, + ) -> None: + from winml.modelkit.commands.export import export + from winml.modelkit.export.cgc import CGCExporter + + output_path = tmp_path / "clip.mlir" + blocked_output = output_path.with_stem("clip_text-encoder") + blocked_output.with_name(f"{blocked_output.name}.data").write_bytes(b"existing") + + with ( + patch( + "winml.modelkit.loader.resolution.resolve_composite_components", + return_value={ + "image-encoder": "image-feature-extraction", + "text-encoder": "feature-extraction", + }, + ), + patch.object(CGCExporter, "export_pytorch") as export_pytorch, + ): + result = runner.invoke( + export, + [ + "--model", + "openai/clip-vit-base-patch32", + "--output", + str(output_path), + "--target", + "cgir", + "--options", + "external-weights=true", + ], + obj={"debug": False}, + ) + + assert result.exit_code != 0 + assert "Output sidecar" in result.output + assert "clip_text-encoder.mlir.data" in result.output + export_pytorch.assert_not_called() + + class TestExportSubmodel: """Test --submodel filters composite export to a single sub-model.""" diff --git a/tests/unit/commands/test_optimize_cli.py b/tests/unit/commands/test_optimize_cli.py index cab204edd..17d9a39af 100644 --- a/tests/unit/commands/test_optimize_cli.py +++ b/tests/unit/commands/test_optimize_cli.py @@ -78,6 +78,7 @@ def test_help_shows_required_flags(self, runner: CliRunner) -> None: "-m", "--output", "-o", + "--disable-ort-graph-optimization", "--ep", "--device", "-d", @@ -268,6 +269,45 @@ def test_device_target_forwarded_to_optimizer( is resolved_ep_device ) + def test_disable_graph_optimization_without_enabling_rewrites( + self, runner: CliRunner, tmp_path: Path + ) -> None: + model_file = tmp_path / "model.onnx" + model_file.touch() + mock_model = _make_mock_model() + + with ( + patch(_LOAD_ONNX, return_value=mock_model), + patch(_SAVE_ONNX), + patch(_OPTIMIZER) as mock_opt_cls, + ): + mock_opt_cls.return_value.optimize.return_value = mock_model + result = runner.invoke( + optimize, + ["-m", str(model_file), "--disable-ort-graph-optimization"], + ) + + assert result.exit_code == 0, result.output + kwargs = mock_opt_cls.return_value.optimize.call_args.kwargs + assert kwargs["ort_graph_optimization"] is False + assert "backend" not in kwargs + assert kwargs["omit_empty_resize_inputs"] is False + + def test_backend_option_removed( + self, runner: CliRunner, tmp_path: Path + ) -> None: + model_file = tmp_path / "model.onnx" + model_file.touch() + + result = runner.invoke( + optimize, + ["-m", str(model_file), "--backend", "cgc", "--device", "gpu"], + ) + + assert result.exit_code != 0 + assert "No such option" in result.output + assert "--backend" in result.output + # ============================================================================= # --check-optim TESTS diff --git a/tests/unit/commands/test_perf_cli.py b/tests/unit/commands/test_perf_cli.py index c40b3509c..9c206c821 100644 --- a/tests/unit/commands/test_perf_cli.py +++ b/tests/unit/commands/test_perf_cli.py @@ -41,6 +41,36 @@ from winml.modelkit.utils.console import SafeConsole +class TestResolveRuntime: + def test_auto_selects_winml_runtime_for_mlir(self) -> None: + assert perf_module._resolve_runtime("auto", "model.mlir") == "winml-runtime" + + def test_explicit_runtime_is_preserved_for_mlir(self) -> None: + assert perf_module._resolve_runtime("winml-ort", "model.mlir") == "winml-ort" + + def test_backend_rejected_for_other_runtime(self, runner: CliRunner) -> None: + result = runner.invoke( + perf, + ["-m", "model.onnx", "--runtime", "winml-ort", "--backend", "cgc"], + obj={}, + ) + + assert result.exit_code == 2 + assert "--backend is only supported with --runtime winml-runtime" in result.output + + def test_ort_backend_rejected_for_mlir(self, runner: CliRunner) -> None: + with runner.isolated_filesystem(): + Path("model.mlir").touch() + result = runner.invoke( + perf, + ["-m", "model.mlir", "--runtime", "winml-runtime", "--backend", "ort"], + obj={}, + ) + + assert result.exit_code == 2 + assert "MLIR inputs require the CGC backend" in result.output + + class TestPerfCacheOptions: @staticmethod def _capture_config( @@ -108,6 +138,26 @@ def test_canonical_flags_reach_benchmark_config( assert config.use_cache is use_cache assert config.rebuild is rebuild + @pytest.mark.parametrize( + ("extra_args", "expected"), + [ + (["--runtime", "winml-runtime"], "cgc"), + (["--runtime", "winml-runtime", "--backend", "ort"], "ort"), + ], + ) + def test_backend_reaches_benchmark_config( + self, + runner: CliRunner, + tmp_path: Path, + extra_args: list[str], + expected: str, + ) -> None: + result, config = self._capture_config(runner, tmp_path, extra_args) + + assert result.exit_code == 0, result.output + assert config is not None + assert config.backend == expected + @pytest.fixture(autouse=True) def mock_resolve_device(): @@ -594,7 +644,8 @@ def test_resolve_device_ep_filters_native_warnings_and_preserves_errors( fake_ep_device.device.ep_name = "QNNExecutionProvider" fake_ep_device.device.device_type = "NPU" - def fake_resolve_device(target: object) -> object: + def fake_resolve_device(target: object, *, backend: str | None = None) -> object: + assert backend is None os.write(2, b"2026 [W:custom-native:, file.cc:1 Probe] hidden warning\n") return target @@ -621,6 +672,95 @@ def registry_instance() -> FakeRegistry: assert "hidden warning" not in stderr assert "useful error" in stderr + def test_runtime_loads_before_ep_registration( + self, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + from winml.modelkit import session as session_module + from winml.modelkit.session.ep_device import resolve_device + + calls: list[str] = [] + fake_ep_device = MagicMock() + + def import_runtime() -> None: + calls.append("runtime") + + class FakeRegistry: + def auto_device(self, _target: object) -> object: + calls.append("auto_device") + return fake_ep_device + + with monkeypatch.context() as local_patch: + local_patch.setattr( + "winml.modelkit.session._runtime_import.import_runtime", + import_runtime, + ) + local_patch.setattr(session_module, "resolve_device", resolve_device) + local_patch.setattr( + session_module.WinMLEPRegistry, + "instance", + staticmethod(FakeRegistry), + ) + + benchmark = PerfBenchmark( + BenchmarkConfig( + model_id="model.onnx", + runtime="winml-runtime", + ep="openvino", + device="gpu", + ) + ) + benchmark._resolve_device_ep() + + assert calls == ["runtime", "auto_device"] + + def test_mlir_resolves_gpu_from_winmlcg_inventory( + self, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + from winml.modelkit import session as session_module + + requested_targets: list[tuple[object, str | None]] = [] + fake_ep_device = MagicMock() + + def resolve_device(target: object, *, backend: str | None = None) -> object: + requested_targets.append((target, backend)) + return SimpleNamespace(ep="winmlcg", device="gpu", source=target.source) + + class FakeRegistry: + def auto_device(self, _target: object) -> object: + return fake_ep_device + + with monkeypatch.context() as local_patch: + local_patch.setattr( + "winml.modelkit.session.runtime_session.import_runtime", + lambda: None, + ) + local_patch.setattr(session_module, "resolve_device", resolve_device) + local_patch.setattr( + session_module.WinMLEPRegistry, + "instance", + staticmethod(FakeRegistry), + ) + + benchmark = PerfBenchmark( + BenchmarkConfig( + model_id="model.mlir", + runtime="winml-runtime", + ) + ) + benchmark._resolve_device_ep() + + requested, backend = requested_targets[0] + assert requested.ep == "auto" + assert requested.device == "auto" + assert requested.source is None + assert backend == "cgc" + assert benchmark._ep_device is fake_ep_device + assert benchmark.config.ep is None + assert benchmark.config.device == "auto" + assert benchmark.resolved_device == "gpu" + def test_onnx_load_model_calls_from_onnx(self, tmp_path: Path) -> None: """ONNX file input should use WinMLAutoModel.from_onnx in _load_model.""" onnx_file = tmp_path / "model.onnx" @@ -834,6 +974,49 @@ def test_cli_onnx_routes_through_perf_benchmark( assert result.exit_code == 0, result.output mock_perf_cls.assert_called_once() + def test_cli_onnx_winml_runtime_preserves_explicit_target( + self, runner: CliRunner, tmp_path: Path + ) -> None: + onnx_file = tmp_path / "model.onnx" + onnx_file.write_bytes(b"fake onnx") + captured: dict[str, BenchmarkConfig] = {} + + def capture_config(config: BenchmarkConfig) -> MagicMock: + captured["config"] = config + mock = MagicMock() + mock.run.return_value = MagicMock() + return mock + + with ( + patch( + "winml.modelkit.commands.perf.PerfBenchmark", + side_effect=capture_config, + ), + patch("winml.modelkit.commands.perf.display_console_report"), + patch("winml.modelkit.commands.perf.write_json_report"), + ): + result = runner.invoke( + perf, + [ + "-m", + str(onnx_file), + "--runtime", + "winml-runtime", + "--ep", + "openvino", + "--device", + "gpu", + "-o", + str(tmp_path / "out.json"), + ], + obj={}, + ) + + assert result.exit_code == 0, result.output + assert captured["config"].ep == "openvino" + assert captured["config"].device == "gpu" + assert "--ep and --ep-options are ignored" not in result.output + def test_cli_onnx_preserves_shape_config(self, runner: CliRunner, tmp_path: Path) -> None: """ONNX input with --shape-config keeps the override for dummy inputs. @@ -1751,6 +1934,13 @@ def test_to_dict_includes_schema_version_and_runtime(self) -> None: assert d["schema_version"] == 2 assert d["benchmark_info"]["runtime"] == "winml-ort" + def test_to_dict_reports_configured_runtime(self) -> None: + config = BenchmarkConfig(model_id="m.mlir", runtime="winml-runtime") + result = BenchmarkResult(config=config) + + info = result.to_dict()["benchmark_info"] + assert info["runtime"] == "winml-runtime" + def test_iterations_reports_configured_count_without_duration(self) -> None: """Without --duration, benchmark_info.iterations is the configured value.""" config = BenchmarkConfig(model_id="m", iterations=100) diff --git a/tests/unit/commands/test_perf_genai.py b/tests/unit/commands/test_perf_genai.py index 626cb361c..bee4593d6 100644 --- a/tests/unit/commands/test_perf_genai.py +++ b/tests/unit/commands/test_perf_genai.py @@ -1826,7 +1826,7 @@ def test_autobuild_without_recipe_rejected( def test_runtime_help_shows_auto_default(self, runner: CliRunner, capture_run: dict) -> None: result = runner.invoke(perf, ["--help"]) assert result.exit_code == 0 - assert "[auto|winml-ort|ort-genai]" in result.output + assert "[auto|winml-ort|ort-genai|winml-runtime]" in result.output assert "default: auto" in result.output assert "config" not in capture_run diff --git a/tests/unit/compiler/test_compiler_configs.py b/tests/unit/compiler/test_compiler_configs.py index 2effc604b..7cdcb8773 100644 --- a/tests/unit/compiler/test_compiler_configs.py +++ b/tests/unit/compiler/test_compiler_configs.py @@ -103,6 +103,13 @@ def test_for_openvino(self): assert config.ep_config.provider == "openvino" assert config.ep_config.enable_ep_context is True + def test_for_winmlcg(self): + """Test WinML Compute Graph factory method.""" + config = WinMLCompileConfig.for_winmlcg() + assert config.ep_config.provider == "winmlcg" + assert config.ep_config.enable_ep_context is True + assert config.ep_config.device == "gpu" + def test_for_vitisai(self, tmp_path, monkeypatch): """Test Vitis AI factory method.""" monkeypatch.setenv("WINML_CACHE_DIR", str(tmp_path)) @@ -277,6 +284,7 @@ class TestForProvider: ("openvino", "openvino"), ("vitisai", "vitisai"), ("nv_tensorrt_rtx", "nvtensorrtrtx"), + ("winmlcg", "winmlcg"), # EPs with enable_ep_context=False → no offline compile step → None ("dml", None), ("cpu", None), @@ -325,6 +333,7 @@ def test_for_provider_custom_ep_returns_none(self): ("nv_tensorrt_rtx", "nvtensorrtrtx"), ("openvino", "openvino"), ("vitisai", "vitisai"), + ("winmlcg", "winmlcg"), ("migraphx", None), ], ) diff --git a/tests/unit/config/test_build.py b/tests/unit/config/test_build.py index e90b5b91c..dca5f5abb 100644 --- a/tests/unit/config/test_build.py +++ b/tests/unit/config/test_build.py @@ -221,7 +221,7 @@ def test_generated_hf_config_can_split_build_and_export_policy_targets( ) monkeypatch.setattr( "winml.modelkit.config.build._apply_target_policy", - lambda config, *, device, precision, ep: target_policy_calls.append( + lambda config, *, device, precision, ep, backend: target_policy_calls.append( (device, precision, ep) ), ) diff --git a/tests/unit/config/test_precision.py b/tests/unit/config/test_precision.py index 85e4f813b..9d8db9ed6 100644 --- a/tests/unit/config/test_precision.py +++ b/tests/unit/config/test_precision.py @@ -246,7 +246,9 @@ def test_all_valid_eps(self) -> None: for ep_name in VALID_EPS: policy = resolve_precision(ep=ep_name) - assert policy.compile_provider == (None if ep_name == "cpu" else ep_name) + assert policy.compile_provider == ( + None if ep_name in ("cpu", "winmlcg") else ep_name + ) def test_ep_accepts_aliases(self) -> None: """resolve_precision should accept shorthand aliases.""" diff --git a/tests/unit/datasets/test_random_dataset.py b/tests/unit/datasets/test_random_dataset.py index fec9e8413..90a3344dc 100644 --- a/tests/unit/datasets/test_random_dataset.py +++ b/tests/unit/datasets/test_random_dataset.py @@ -133,6 +133,52 @@ def test_random_dataset_with_model_path(self, simple_onnx_model: Path) -> None: assert "A" in sample # Input name from ONNX model assert sample["A"].shape == (1, 4) + def test_random_dataset_with_io_config(self) -> None: + from unittest.mock import patch + + from winml.modelkit.datasets import RandomDataset + + io_config = { + "input_names": ["input"], + "input_shapes": [[1, 3]], + "input_types": [np.dtype("float32")], + "value_ranges": {"input": (-1.0, 1.0)}, + } + with patch( + "winml.modelkit.onnx.get_io_config", + side_effect=AssertionError("artifact must not be parsed"), + ): + dataset = RandomDataset(model_path=None, io_config=io_config, max_samples=2) + + assert len(dataset) == 2 + assert dataset[0]["input"].shape == (1, 3) + + def test_model_path_remains_authoritative(self, simple_onnx_model: Path) -> None: + from unittest.mock import patch + + from winml.modelkit.datasets import RandomDataset + + conflicting_io_config = { + "input_names": ["wrong_input"], + "input_shapes": [[1, 99]], + "input_types": [np.dtype("int64")], + } + with patch("winml.modelkit.onnx.get_io_config") as get_io_config: + get_io_config.return_value = { + "input_names": ["A"], + "input_shapes": [[1, 4]], + "input_types": [np.dtype("float32")], + } + dataset = RandomDataset( + model_path=str(simple_onnx_model), + io_config=conflicting_io_config, + max_samples=1, + ) + + get_io_config.assert_called_once_with(str(simple_onnx_model)) + assert set(dataset[0]) == {"A", "sample_id"} + assert dataset[0]["A"].shape == (1, 4) + def test_random_dataset_generates_correct_dtype(self, simple_onnx_model: Path) -> None: """RandomDataset should generate data with correct dtype.""" import torch diff --git a/tests/unit/ep_path/test_ep_path.py b/tests/unit/ep_path/test_ep_path.py index 85a06e41c..c6b8cc1bb 100644 --- a/tests/unit/ep_path/test_ep_path.py +++ b/tests/unit/ep_path/test_ep_path.py @@ -131,6 +131,20 @@ def test_ep_catalog_has_five_plugin_eps(self) -> None: "NvTensorRTRTXExecutionProvider", } + def test_winmlcg_uses_pypi_source(self) -> None: + sources = [ + source + for source in _default_ep_sources() + if "WinMLCGExecutionProvider" in source.iter_eps() + ] + assert sources == [ + PyPISource( + distribution="windowsml", + relative_dll="windowsml/lib/WinMLCGEp.dll", + eps=("WinMLCGExecutionProvider",), + ) + ] + def test_ep_catalog_uses_canonical_casing_for_nvidia(self) -> None: assert EP_CATALOG.dll_name_for("NvTensorRTRTXExecutionProvider") is not None assert EP_CATALOG.dll_name_for("NvTensorRtRtxExecutionProvider") is None diff --git a/tests/unit/eval/test_eval.py b/tests/unit/eval/test_eval.py index 7021cef34..2df642ae3 100644 --- a/tests/unit/eval/test_eval.py +++ b/tests/unit/eval/test_eval.py @@ -81,6 +81,61 @@ def test_config_roundtrip(self): assert restored.dataset.path == config.dataset.path assert restored.dataset.columns_mapping == config.dataset.columns_mapping + +class TestPrintConfig: + def test_compare_shows_candidate_and_onnx_reference_environments(self) -> None: + import importlib + + from rich.console import Console + + eval_mod = importlib.import_module("winml.modelkit.eval.evaluate") + + console = Console(record=True, width=120) + config = WinMLEvaluationConfig( + model_path="candidate.mlir", + reference_path="reference.onnx", + runtime="winml-runtime", + device="gpu", + ep="winmlcg", + reference_device="gpu", + reference_ep="dml", + mode="compare", + ) + + with patch.object(eval_mod, "Console", return_value=console): + eval_mod.print_config(config) + + text = console.export_text() + assert "Candidate: candidate.mlir" in text + assert "Candidate runtime: winml-runtime" in text + assert "Candidate device: gpu" in text + assert "Candidate EP:" not in text + assert "Reference: reference.onnx" in text + assert "Reference runtime: winml-ort" in text + assert "Reference device: gpu" in text + assert "Reference EP: dml" in text + + def test_compare_runtime_ort_shows_candidate_ep(self) -> None: + import importlib + + from rich.console import Console + + eval_mod = importlib.import_module("winml.modelkit.eval.evaluate") + console = Console(record=True, width=120) + config = WinMLEvaluationConfig( + model_path="candidate.onnx", + reference_path="reference.onnx", + runtime="winml-runtime", + backend="ort", + ep="dml", + mode="compare", + ) + + with patch.object(eval_mod, "Console", return_value=console): + eval_mod.print_config(config) + + assert "Candidate EP: dml" in console.export_text() + def test_config_roundtrip_preserves_revision(self): """DatasetConfig.revision survives to_dict/from_dict roundtrip.""" config = WinMLEvaluationConfig( @@ -94,6 +149,17 @@ def test_config_roundtrip_preserves_revision(self): restored = WinMLEvaluationConfig.from_dict(config.to_dict()) assert restored.dataset.revision == "refs/convert/parquet" + def test_config_roundtrip_preserves_runtime_backend(self): + config = WinMLEvaluationConfig( + model_path="model.onnx", + runtime="winml-runtime", + backend="ort", + ) + + restored = WinMLEvaluationConfig.from_dict(config.to_dict()) + + assert restored.backend == "ort" + def test_dataset_config_revision_default_is_none(self): """Revision defaults to None when not specified.""" ds = DatasetConfig(path="some-dataset") @@ -119,6 +185,7 @@ def test_config_roundtrip_preserves_input_data(self): def test_config_roundtrip_preserves_cache_controls(self): config = WinMLEvaluationConfig( model_id="test/model", + runtime="winml-runtime", use_cache=False, rebuild=True, ) @@ -411,7 +478,16 @@ def test_none_mode_normalizes_to_onnx(self): result = eval_mod.evaluate(config) assert result.config.mode == "onnx" - def test_onnx_compare_ignores_explicit_task_and_skips_resolution(self): + @pytest.mark.parametrize( + "model_path,runtime,task", + [ + ("cand.onnx", "winml-ort", "image-classification"), + ("cand.mlir", "winml-runtime", "text-generation"), + ], + ) + def test_onnx_compare_ignores_explicit_task_and_skips_resolution( + self, model_path, runtime, task, + ): """Two-ONNX compare preserves all raw outputs by clearing the task.""" import importlib import sys @@ -421,10 +497,11 @@ def test_onnx_compare_ignores_explicit_task_and_skips_resolution(self): ) or importlib.import_module("winml.modelkit.eval.evaluate") config = WinMLEvaluationConfig( - model_path="cand.onnx", + model_path=model_path, + runtime=runtime, reference_path="ref.onnx", mode="compare", - task="image-classification", + task=task, ) evaluator = MagicMock() @@ -452,6 +529,19 @@ def test_onnx_compare_ignores_explicit_task_and_skips_resolution(self): load_model.assert_called_once_with(result.config) evaluator_factory.assert_called_once_with(result.config, candidate) + @pytest.mark.parametrize("suffix", [".mlir", ".MLIR"]) + def test_from_mlir_rejects_text_generation_before_wrapper_creation(self, suffix): + from winml.modelkit.models import WinMLAutoModel + + with ( + patch("winml.modelkit.models.auto.get_winml_class") as get_winml_class, + pytest.raises(ValueError, match="from_mlir does not support task='text-generation'"), + ): + WinMLAutoModel.from_mlir( + f"model{suffix}", ep_device=MagicMock(), task="text-generation", + ) + get_winml_class.assert_not_called() + def test_no_dataset_no_default_raises(self): """Tasks without a default dataset raise ValueError.""" import importlib @@ -1650,6 +1740,8 @@ def test_load_onnx_without_model_id_returns_generic_winml_model(self): model_path="candidate.onnx", reference_path="reference.onnx", mode="compare", + runtime="winml-runtime", + backend="cgc", device="cpu", ) @@ -1663,8 +1755,42 @@ def test_load_onnx_without_model_id_returns_generic_winml_model(self): mock_auto.from_onnx.assert_called_once() assert mock_auto.from_onnx.call_args.kwargs["hf_config"] is None assert mock_auto.from_onnx.call_args.kwargs["task"] is None + assert mock_auto.from_onnx.call_args.kwargs["runtime"] == "winml-runtime" assert mock_auto.from_onnx.call_args.kwargs["skip_build"] is True + def test_load_mlir_uses_runtime_model_loader(self): + import importlib + import sys + + eval_mod = sys.modules.get( + "winml.modelkit.eval.evaluate", + ) or importlib.import_module("winml.modelkit.eval.evaluate") + + mock_model = MagicMock() + mock_auto = MagicMock() + mock_auto.from_mlir.return_value = mock_model + config = WinMLEvaluationConfig( + model_path="candidate.mlir", + runtime="winml-runtime", + task="text-generation", + device="gpu", + ) + + with ( + patch.dict( + "sys.modules", + {"winml.modelkit.models": MagicMock(WinMLAutoModel=mock_auto)}, + ), + patch("winml.modelkit.session.runtime_session.import_runtime"), + ): + result = eval_mod.load_model(config) + + assert result is mock_model + mock_auto.from_mlir.assert_called_once() + assert mock_auto.from_mlir.call_args.kwargs["mlir_path"] == "candidate.mlir" + assert mock_auto.from_mlir.call_args.kwargs["task"] == "text-generation" + assert mock_auto.from_mlir.call_args.kwargs["runtime"] == "winml-runtime" + def test_make_onnx_reference_config_uses_independent_environment(self): from winml.modelkit.eval.tensor_similarity_evaluator import ( _make_reference_config, @@ -1678,6 +1804,8 @@ def test_make_onnx_reference_config_uses_independent_environment(self): device_luid="0x00000000_0x00000001", reference_device_luid="0x00000000_0x00000002", mode="compare", + runtime="winml-runtime", + backend="cgc", ) reference = _make_reference_config(config) @@ -1686,6 +1814,7 @@ def test_make_onnx_reference_config_uses_independent_environment(self): assert reference.model_id is None assert reference.reference_path is None assert reference.runtime == "winml-ort" + assert reference.backend is None assert reference.device == "gpu" assert reference.device_luid == "0x00000000_0x00000002" assert reference.ep == "dml" @@ -1702,6 +1831,8 @@ def test_make_default_hf_reference_config_uses_native_defaults(self): task="image-classification", device_luid="0x00000000_0x00000001", mode="compare", + runtime="winml-runtime", + backend="cgc", ) reference = _make_reference_config(config) @@ -1710,6 +1841,7 @@ def test_make_default_hf_reference_config_uses_native_defaults(self): assert reference.model_path is None assert reference.reference_path is None assert reference.runtime == "pytorch" + assert reference.backend is None assert reference.device == "cpu" assert reference.device_luid is None assert reference.ep is None @@ -1747,7 +1879,12 @@ def test_auto_target_retries_cpu_after_ort_runtime_failure(self, caplog): ) config._auto_device_selected = True - def resolve_target(target: EPDeviceTarget) -> EPDeviceTarget: + def resolve_target( + target: EPDeviceTarget, + *, + backend: str | None = None, + ) -> EPDeviceTarget: + assert backend is None if target.device == "gpu": return EPDeviceTarget(ep="DmlExecutionProvider", device=target.device) return EPDeviceTarget(ep="CPUExecutionProvider", device="cpu") diff --git a/tests/unit/eval/test_tensor_similarity_evaluator.py b/tests/unit/eval/test_tensor_similarity_evaluator.py index 0a39a0acd..3c2887668 100644 --- a/tests/unit/eval/test_tensor_similarity_evaluator.py +++ b/tests/unit/eval/test_tensor_similarity_evaluator.py @@ -213,8 +213,8 @@ def test_loads_onnx_reference_with_independent_config(self, monkeypatch): assert reference_config.ep == "dml" assert reference_config.task is None assert load_model.call_args.kwargs["torch_dtype"] is torch.float32 - # RandomDataset is built over the candidate ONNX I/O. - assert evaluator.data.kwargs["model_path"].endswith("cand.onnx") + # RandomDataset consumes the already-loaded candidate model schema. + assert evaluator.data.kwargs["io_config"] is candidate.io_config assert evaluator.data.kwargs["max_samples"] == 5 assert evaluator.data.kwargs["seed"] == 1 diff --git a/tests/unit/export/cgc/test_exporter.py b/tests/unit/export/cgc/test_exporter.py new file mode 100644 index 000000000..8f4b597ce --- /dev/null +++ b/tests/unit/export/cgc/test_exporter.py @@ -0,0 +1,395 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- + +from pathlib import Path +from unittest.mock import patch + +import numpy as np +import pytest +from onnx import ModelProto, TensorProto, helper, numpy_helper, save_model + +from winml.modelkit.export import WinMLExportConfig +from winml.modelkit.export.cgc import CGCExporter, CGCExportResult, CGCOptions +from winml.modelkit.export.cgc.foundry import FoundryCompileError + + +def _make_model() -> ModelProto: + input_info = helper.make_tensor_value_info("input", TensorProto.FLOAT, [1, 2]) + output_info = helper.make_tensor_value_info("output", TensorProto.FLOAT, [1, 2]) + weight = numpy_helper.from_array( + np.array([1.0, 2.0], dtype=np.float32), + name="weight", + ) + bias = helper.make_tensor( + "bias", + TensorProto.FLOAT, + [2], + [0.5, 1.5], + ) + nodes = [ + helper.make_node("Add", ["input", "weight"], ["hidden"]), + helper.make_node("Add", ["hidden", "bias"], ["output"]), + ] + graph = helper.make_graph( + nodes, + "test", + [input_info], + [output_info], + initializer=[weight, bias], + ) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + + +def test_pytorch_export_hides_onnx_intermediate(tmp_path: Path) -> None: + output = tmp_path / "model.mlir" + exporter = CGCExporter(CGCOptions()) + intermediate_paths: list[Path] = [] + export_stats = {"nodes": 1} + + def fake_export_onnx(**kwargs): + intermediate_path = Path(kwargs["output_path"]) + intermediate_path.write_bytes(b"onnx") + intermediate_paths.append(intermediate_path) + return export_stats + + with ( + patch( + "winml.modelkit.export.export_pytorch", + side_effect=fake_export_onnx, + ) as export_onnx, + patch.object( + exporter, + "export_onnx", + return_value=CGCExportResult(input_names=(), output_names=()), + ) as export_cgc, + ): + result = exporter.export_pytorch( + model=object(), + output_path=output, + export_config=WinMLExportConfig(), + model_id="model-id", + task="feature-extraction", + verbose=True, + enable_reporting=False, + ) + + assert result.export_stats is export_stats + assert result.input_names == () + assert result.output_names == () + export_onnx.assert_called_once() + export_cgc.assert_called_once_with( + model=intermediate_paths[0], + output_path=output, + ) + assert intermediate_paths[0].name == "model.onnx" + assert not intermediate_paths[0].parent.exists() + + +class _ExternalMlirCompiler: + def __enter__(self): + return self + + def __exit__(self, *_args): + return None + + def compile_onnx(self, _source, **kwargs): + kwargs["output_data_file"].write_bytes(b"new weights") + return b"module { cgc.test }" + + +def test_external_mlir_overwrite_replaces_sidecar(tmp_path: Path) -> None: + source = tmp_path / "source.onnx" + output = tmp_path / "model.mlir" + sidecar = output.with_name(f"{output.name}.data") + metadata = tmp_path / "model_metadata.json" + save_model(_make_model(), str(source)) + sidecar.write_bytes(b"old weights") + + with patch( + "winml.modelkit.export.cgc.exporter.FoundryCompiler", + return_value=_ExternalMlirCompiler(), + ): + CGCExporter(CGCOptions(external_weights=True)).export( + source, + output, + ) + + assert output.read_text(encoding="utf-8") == "module { cgc.test }" + assert sidecar.read_bytes() == b"new weights" + assert metadata.is_file() + + +def test_external_mlir_failure_preserves_existing_bundle(tmp_path: Path) -> None: + source = tmp_path / "source.onnx" + output = tmp_path / "model.mlir" + sidecar = output.with_name(f"{output.name}.data") + metadata = tmp_path / "model_metadata.json" + save_model(_make_model(), str(source)) + output.write_text("old mlir", encoding="utf-8") + sidecar.write_bytes(b"old weights") + metadata.write_text("old metadata", encoding="utf-8") + exporter = CGCExporter(CGCOptions(external_weights=True)) + + with ( + patch( + "winml.modelkit.export.cgc.exporter.FoundryCompiler", + return_value=_ExternalMlirCompiler(), + ), + patch.object( + exporter, + "_write_io_metadata", + side_effect=RuntimeError("metadata write failed"), + ), + pytest.raises(RuntimeError, match="metadata write failed"), + ): + exporter.export(source, output) + + assert output.read_text(encoding="utf-8") == "old mlir" + assert sidecar.read_bytes() == b"old weights" + assert metadata.read_text(encoding="utf-8") == "old metadata" + + +@pytest.mark.parametrize("existing", [False, True]) +@pytest.mark.parametrize("weights_mode", ["external", "absent", "embedded"]) +@pytest.mark.parametrize("failure", [None, "backup", "weights", "metadata", "model"]) +def test_mlir_bundle_publication( + tmp_path: Path, monkeypatch, existing: bool, weights_mode: str, failure: str | None, +) -> None: + source = tmp_path / "source.onnx" + model = _make_model() + save_model(model, source) + original_source = source.read_bytes() + output = tmp_path / "model.mlir" + exporter = CGCExporter(CGCOptions(external_weights=weights_mode != "embedded")) + + class Compiler(_ExternalMlirCompiler): + def compile_onnx(self, serialized, **kwargs): + data_path = kwargs["output_data_file"] + if data_path is not None and weights_mode == "external": + data_path.write_bytes(serialized) + return b"module { cgc.test }" + + monkeypatch.setattr( + "winml.modelkit.export.cgc.exporter.FoundryCompiler", Compiler, + ) + if existing: + exporter.export_onnx(source, output) + if weights_mode == "absent": + output.with_name(f"{output.name}.data").write_bytes(original_source) + artifacts = exporter.output_artifacts(output) + before = {path.name: path.read_bytes() for path in artifacts if path.exists()} + model.graph.input[0].name = "updated_input" + model.graph.node[0].input[0] = "updated_input" + save_model(model, source) + updated_source = source.read_bytes() + expected_dir = tmp_path / "expected" + exporter.export_onnx(source, expected_dir / output.name) + expected = { + path.name: path.read_bytes() + for path in exporter.output_artifacts(expected_dir / output.name) + if path.exists() + } + real_replace = Path.replace + failure_names = { + "weights": "model.mlir.data", "metadata": "model_metadata.json", "model": "model.mlir", + } + triggered = False + + def fail_once(path, target): + nonlocal triggered + target = Path(target) + is_backup = path.parent == tmp_path and target.parent.name == "backup" + is_publication = target.parent == tmp_path and path.name == target.name + if not triggered and ( + (failure == "backup" and is_backup and path.name == output.name) + or (is_publication and target.name == failure_names.get(failure)) + ): + triggered = True + raise PermissionError("publication probe") + return real_replace(path, target) + + should_fail = ( + failure in ("metadata", "model") + or (failure == "backup" and existing) + or (failure == "weights" and weights_mode == "external") + ) + with monkeypatch.context() as scoped: + scoped.setattr(Path, "replace", fail_once) + if should_fail: + with pytest.raises(PermissionError, match="publication probe"): + exporter.export_onnx(source, output) + else: + exporter.export_onnx(source, output) + after = {path.name: path.read_bytes() for path in artifacts if path.exists()} + assert triggered == should_fail + assert after == (before if should_fail else expected) + assert source.read_bytes() == updated_source + assert not list(tmp_path.glob(".model.mlir.*")) + + +def test_export_prints_progress(capsys: pytest.CaptureFixture[str], tmp_path: Path) -> None: + source = tmp_path / "source.onnx" + output = tmp_path / "model.mlir" + save_model(_make_model(), str(source)) + + class FakeCompiler: + def __enter__(self): + return self + + def __exit__(self, *_args): + return None + + def compile_onnx(self, _source, **_kwargs): + return b"module { cgc.test }" + + with patch( + "winml.modelkit.export.cgc.exporter.FoundryCompiler", + return_value=FakeCompiler(), + ): + CGCExporter(CGCOptions()).export(source, output) + + console_output = capsys.readouterr().out + assert "ONNX TO CGC EXPORT PROCESS" in console_output + assert "Input:" in console_output + assert "Output:" in console_output + assert "Format: CGC MLIR" in console_output + assert "CGC EXPORT COMPLETE" in console_output + + +@pytest.mark.parametrize( + ("options", "expected_names"), + [ + (CGCOptions(), ("model.mlir", "model_metadata.json")), + ( + CGCOptions(external_weights=True), + ("model.mlir", "model.mlir.data", "model_metadata.json"), + ), + ], +) +def test_output_artifacts_depend_only_on_external_weights( + tmp_path: Path, + options: CGCOptions, + expected_names: tuple[str, ...], +) -> None: + artifacts = CGCExporter(options).output_artifacts(tmp_path / "model.mlir") + assert tuple(path.name for path in artifacts) == expected_names + + +@pytest.mark.parametrize("dimension", ["batch_size", "batch", "batchSize", 1, 4]) +@pytest.mark.parametrize("explicit", ["", "batch_size=3", "seq=128"]) +def test_auto_freeze_matches_input_symbol_and_preserves_options(tmp_path, dimension, explicit): + model = _make_model() + model.graph.input[0].CopyFrom( + helper.make_tensor_value_info("input", TensorProto.FLOAT, [dimension, 2]), + ) + source = tmp_path / "source.onnx" + save_model(model, source) + exporter = CGCExporter(CGCOptions(freeze_dims=explicit)) + with patch.object(exporter, "_export_mlir"): + exporter.export_onnx(source, tmp_path / "model.mlir") + if explicit: + name, size = explicit.split("=") + expected = {name: int(size)} + if name != "batch_size" and dimension == "batch_size": + expected["batch_size"] = 1 + assert exporter._freeze_dims == expected + else: + assert exporter._freeze_dims == ({"batch_size": 1} if dimension == "batch_size" else {}) + save_model(_make_model(), source) + exporter.export_onnx(source, tmp_path / "static.mlir") + if explicit: + name, size = explicit.split("=") + assert exporter._freeze_dims == {name: int(size)} + else: + assert exporter._freeze_dims == {} + + +@pytest.mark.parametrize("location", ["output", "value_info", "initializer"]) +def test_auto_freeze_ignores_non_input_symbols(tmp_path, location): + model = _make_model() + value = helper.make_tensor_value_info("weight", TensorProto.FLOAT, ["batch_size"]) + if location == "initializer": + model.graph.input.append(value) + else: + getattr(model.graph, location).append(value) + source = tmp_path / "source.onnx" + save_model(model, source) + exporter = CGCExporter(CGCOptions()) + with patch.object(exporter, "_export_mlir"): + exporter.export_onnx(source, tmp_path / "model.mlir") + assert exporter._freeze_dims == {} + + +def test_freeze_dims_are_forwarded_to_foundry(tmp_path: Path) -> None: + source = tmp_path / "source.onnx" + output = tmp_path / "model.mlir" + save_model(_make_model(), str(source)) + + class FakeCompiler: + def __enter__(self): + return self + + def __exit__(self, *_args): + return None + + def compile_onnx(self, _source, **kwargs): + assert kwargs["freeze_dims"] == {"batch": 1, "seq": 128} + return b"module { cgc.test }" + + with patch( + "winml.modelkit.export.cgc.exporter.FoundryCompiler", + return_value=FakeCompiler(), + ): + CGCExporter( + CGCOptions(freeze_dims="batch=1,seq=128") + ).export_onnx(source, output) + + +@pytest.mark.parametrize( + "freeze_dims", + ["batch", "=1", "batch=x", "batch=0", "batch=1,batch=2"], +) +def test_invalid_freeze_dims_are_rejected(freeze_dims: str) -> None: + with pytest.raises(ValueError, match="freeze-dims"): + CGCExporter(CGCOptions(freeze_dims=freeze_dims)) + + +@pytest.mark.parametrize( + ("error", "expected"), + [ + ( + FoundryCompileError( + 9, + native_message="Operation is not registered.", + unsupported_op="com.example.CustomOp", + ), + "ONNX operator 'com.example.CustomOp' is not supported.", + ), + ( + FoundryCompileError( + 11, + native_message="File does not exist.", + missing_external_data="weights/model.data", + ), + "ONNX external weights file was not found: 'weights/model.data'.", + ), + ( + FoundryCompileError( + 3, + native_message="Could not infer output shape.", + ), + "--options freeze-dims=batch=1,seq=128", + ), + ], +) +def test_exporter_formats_foundry_diagnostics( + error: FoundryCompileError, + expected: str, +) -> None: + message = CGCExporter(CGCOptions())._format_foundry_error(error) + + assert f"[{error.result_name}]" in message + assert error.native_message in message + assert expected in message diff --git a/tests/unit/export/cgc/test_foundry.py b/tests/unit/export/cgc/test_foundry.py new file mode 100644 index 000000000..24841dad7 --- /dev/null +++ b/tests/unit/export/cgc/test_foundry.py @@ -0,0 +1,203 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Tests for FoundryToolbox discovery.""" + +from __future__ import annotations + +import ctypes +from importlib import metadata +from typing import TYPE_CHECKING +from unittest.mock import Mock, patch + +import pytest + +from winml.modelkit.export.cgc.foundry import ( + FoundryCompileError, + FoundryCompiler, + FoundryToolboxUnavailableError, + _FdyOverrideDynamicDimsByDimNamePassDescriptor, + find_foundry_toolbox, +) + + +if TYPE_CHECKING: + from pathlib import Path + + +def test_find_foundry_toolbox_in_windowsml_wheel(tmp_path: Path) -> None: + package_dir = tmp_path / "windowsml" + dll_path = package_dir / "lib" / "FoundryToolbox.dll" + dll_path.parent.mkdir(parents=True) + dll_path.write_bytes(b"dll") + distribution = Mock(version="2.6.10.dev0") + distribution.locate_file.return_value = dll_path + + with patch( + "winml.modelkit.export.cgc.foundry.metadata.distribution", + return_value=distribution, + ): + assert find_foundry_toolbox() == dll_path.resolve() + distribution.locate_file.assert_called_once_with( + "windowsml/lib/FoundryToolbox.dll" + ) + + +def test_find_foundry_toolbox_requires_windowsml_distribution() -> None: + with ( + patch( + "winml.modelkit.export.cgc.foundry.metadata.distribution", + side_effect=metadata.PackageNotFoundError, + ), + pytest.raises( + FoundryToolboxUnavailableError, + match="requires a windowsml wheel", + ), + ): + find_foundry_toolbox() + + +def test_find_foundry_toolbox_requires_dll_in_wheel(tmp_path: Path) -> None: + distribution = Mock(version="2.6.8.dev0") + distribution.locate_file.return_value = ( + tmp_path / "windowsml" / "lib" / "FoundryToolbox.dll" + ) + + with ( + patch( + "winml.modelkit.export.cgc.foundry.metadata.distribution", + return_value=distribution, + ), + pytest.raises( + FoundryToolboxUnavailableError, + match=r"windowsml wheel \(2\.6\.8\.dev0\) does not contain", + ), + ): + find_foundry_toolbox() + + +def _compiler_with_diagnostics( + *, + message: bytes, + unsupported_op: bytes | None = None, + missing_external_data: bytes | None = None, +) -> FoundryCompiler: + compiler = object.__new__(FoundryCompiler) + compiler._compiler = object() + compiler._dll = Mock( + FdyCompilerGetLastError=Mock(return_value=message), + FdyCompilerGetLastUnsupportedOpName=Mock(return_value=unsupported_op), + FdyCompilerGetLastMissingExternalDataFile=Mock( + return_value=missing_external_data + ), + ) + return compiler + + +def test_foundry_error_classifies_result_and_preserves_native_message() -> None: + compiler = _compiler_with_diagnostics( + message=b"Could not infer output shape.", + ) + + error = compiler._last_error(3) + + assert isinstance(error, FoundryCompileError) + assert error.result_code == 3 + assert error.result_name == "SHAPE_INFERENCE" + assert error.native_message == "Could not infer output shape." + assert str(error) == ( + "Foundry compiler failed [SHAPE_INFERENCE]: " + "Could not infer output shape." + ) + + +def test_foundry_error_reports_first_unsupported_operator() -> None: + compiler = _compiler_with_diagnostics( + message=b"Operation is not registered.", + unsupported_op=b"com.example.CustomOp", + ) + + error = compiler._last_error(9) + + assert error.unsupported_op == "com.example.CustomOp" + assert error.result_name == "UNSUPPORTED_OP" + assert "com.example.CustomOp" not in str(error) + + +def test_foundry_error_reports_missing_external_data_path() -> None: + compiler = _compiler_with_diagnostics( + message=b"External data file does not exist.", + missing_external_data=b"weights/model.data", + ) + + error = compiler._last_error(11) + + assert error.missing_external_data == "weights/model.data" + assert error.result_name == "MISSING_EXTERNAL_DATA" + assert "weights/model.data" not in str(error) + + +def test_freeze_dims_populate_foundry_pass_descriptor(tmp_path: Path) -> None: + observed: dict[str, object] = {} + payload = b"module { cgc.test }" + + class FakeDLL: + @staticmethod + def FdyCompilerCompile( # noqa: N802 + _compiler, _source, options_pointer, module_pointer + ): + options = options_pointer._obj + assert options.version == 4 + assert not options.safetensorsFiles + assert options.safetensorsFileCount == 0 + assert options.passCount == 1 + descriptor = ctypes.cast( + options.passes[0], + ctypes.POINTER(_FdyOverrideDynamicDimsByDimNamePassDescriptor), + ).contents + observed["kind"] = descriptor.descriptor.kind + assert descriptor.descriptor.stage == 0 + observed["names"] = tuple( + descriptor.names[index].decode("utf-8") + for index in range(descriptor.count) + ) + observed["values"] = tuple( + descriptor.values[index] for index in range(descriptor.count) + ) + module_pointer._obj.value = 1 + return 0 + + @staticmethod + def FdyModuleSerialize( # noqa: N802 + _module, _format, output, size_pointer + ): + size_pointer._obj.value = len(payload) + if output.data: + ctypes.memmove(output.data, payload, len(payload)) + return 0 + + @staticmethod + def FdyModuleDestroy(_module): # noqa: N802 + return None + + compiler = object.__new__(FoundryCompiler) + compiler._compiler = ctypes.c_void_p(1) + compiler._dll = FakeDLL() + + result = compiler.compile_onnx( + b"onnx", + model_directory=tmp_path, + update_opset=True, + topo_sort_nodes=True, + include_initializers=True, + enable_lazy_external_data=False, + freeze_dims={"batch": 1, "seq": 128}, + ) + + assert result == payload + assert observed == { + "kind": 15, + "names": ("batch", "seq"), + "values": (1, 128), + } diff --git a/tests/unit/models/auto/test_auto_onnx.py b/tests/unit/models/auto/test_auto_onnx.py index 71bc94c48..ca1f992d2 100644 --- a/tests/unit/models/auto/test_auto_onnx.py +++ b/tests/unit/models/auto/test_auto_onnx.py @@ -19,6 +19,7 @@ from unittest.mock import MagicMock, patch import pytest +from onnx import TensorProto, helper, save_model from winml.modelkit.ep_path import BuiltinSource, EPEntry from winml.modelkit.models.auto import WinMLAutoModel @@ -36,9 +37,18 @@ def cpu_ep_device(): @pytest.fixture() def fake_onnx(tmp_path: Path) -> Path: - """Create a fake ONNX file for testing.""" + """Create a minimal valid ONNX file for testing.""" onnx_file = tmp_path / "model.onnx" - onnx_file.write_bytes(b"fake-onnx") + graph = helper.make_graph( + [helper.make_node("Identity", ["input"], ["output"])], + "test", + [helper.make_tensor_value_info("input", TensorProto.FLOAT, [1])], + [helper.make_tensor_value_info("output", TensorProto.FLOAT, [1])], + ) + save_model( + helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]), + onnx_file, + ) return onnx_file @@ -73,6 +83,63 @@ def _make_cpu_ep_device_with_bridge_name() -> WinMLEPDevice: class TestFromOnnx: """Test WinMLAutoModel.from_onnx().""" + def test_winml_runtime_passes_onnx_directly_to_runtime( + self, fake_onnx: Path, cpu_ep_device: EPDeviceTarget + ) -> None: + wrapper = MagicMock() + wrapper_kwargs = {} + + def create_wrapper(**kwargs): + wrapper_kwargs.update(kwargs) + return wrapper + + with ( + patch("winml.modelkit.build.build_onnx_model") as build, + patch( + "winml.modelkit.models.auto.get_winml_class", + return_value=create_wrapper, + ), + ): + model = WinMLAutoModel.from_onnx( + fake_onnx, + ep_device=cpu_ep_device, + task="image-classification", + runtime="winml-runtime", + skip_build=True, + ) + + assert model is wrapper + build.assert_not_called() + assert wrapper_kwargs["onnx_path"] == fake_onnx + assert wrapper_kwargs["runtime"] == "winml-runtime" + + def test_winml_cg_ep_passes_onnx_directly_to_ort( + self, fake_onnx: Path + ) -> None: + ep_device = MagicMock() + ep_device.ep_short_name = "winmlcg" + ep_device.device.device_type = "GPU" + ep_device.device.ep_name = "WinMLCGExecutionProvider" + wrapper = MagicMock() + + with ( + patch("winml.modelkit.build.build_onnx_model") as build, + patch( + "winml.modelkit.models.auto.get_winml_class", + return_value=lambda **kwargs: (wrapper, kwargs), + ), + ): + model, kwargs = WinMLAutoModel.from_onnx( + fake_onnx, + ep_device=ep_device, + task="image-classification", + ) + + build.assert_not_called() + assert model is wrapper + assert kwargs["onnx_path"] == fake_onnx + assert kwargs["runtime"] == "winml-ort" + def test_auto_generates_config_when_none( self, fake_onnx: Path, tmp_path: Path, cpu_ep_device: EPDeviceTarget ): @@ -597,22 +664,21 @@ def test_replacing_same_path_metadata_gets_different_model_dir(self, tmp_path: P def test_onnx_model_hash_includes_external_data_metadata(self, tmp_path: Path): """Changing external data metadata changes the ONNX model hash.""" import numpy as np - import onnx from winml.modelkit.onnx import get_onnx_model_hash onnx_path = tmp_path / "external.onnx" data_path = tmp_path / "external.onnx.data" - tensor = onnx.helper.make_tensor( + tensor = helper.make_tensor( "weight", - onnx.TensorProto.FLOAT, + TensorProto.FLOAT, [4], np.arange(4, dtype=np.float32).tobytes(), raw=True, ) - graph = onnx.helper.make_graph([], "external-data-test", [], [], [tensor]) - model = onnx.helper.make_model(graph) - onnx.save_model( + graph = helper.make_graph([], "external-data-test", [], [], [tensor]) + model = helper.make_model(graph) + save_model( model, str(onnx_path), save_as_external_data=True, diff --git a/tests/unit/models/auto/test_automodel.py b/tests/unit/models/auto/test_automodel.py index 21e997e6f..c036d6fb3 100644 --- a/tests/unit/models/auto/test_automodel.py +++ b/tests/unit/models/auto/test_automodel.py @@ -260,7 +260,7 @@ def test_single_model_passes_runtime_options_factory_to_session(monkeypatch): import winml.modelkit.models.winml.base as base_module session = MagicMock() - monkeypatch.setattr(base_module, "WinMLSession", session) + monkeypatch.setitem(base_module.SESSION_CLASSES, "winml-ort", session) ep_device = MagicMock() session_options = MagicMock() @@ -272,7 +272,7 @@ def test_single_model_passes_runtime_options_factory_to_session(monkeypatch): ) session.assert_called_once_with( - onnx_path=base_module.Path("model.onnx"), + base_module.Path("model.onnx"), ep_device=ep_device, provider_options={"key": "value"}, session_options=session_options, diff --git a/tests/unit/models/auto/test_from_pretrained_ep.py b/tests/unit/models/auto/test_from_pretrained_ep.py index abc86415c..4bd525c4d 100644 --- a/tests/unit/models/auto/test_from_pretrained_ep.py +++ b/tests/unit/models/auto/test_from_pretrained_ep.py @@ -39,6 +39,8 @@ def _install_stubs(monkeypatch: pytest.MonkeyPatch, *, compile_provider: str | N received: dict[str, Any] = {} fake_build_config = MagicMock() + fake_build_config.skip_optimize = False + fake_build_config.quant = MagicMock() if compile_provider is None: fake_build_config.compile = None else: @@ -56,13 +58,18 @@ def generate_config(*_args: Any, **kwargs: Any) -> MagicMock: fake_ep_device = MagicMock() fake_ep_device.device.device_type = "CPU" fake_ep_device.device.ep_name = "CPUExecutionProvider" + + def resolve_target(target: EPDeviceTarget, **kwargs: Any) -> EPDeviceTarget: + received["resolve_backend"] = kwargs.get("backend") + return EPDeviceTarget( + ep=target.ep if target.ep != "auto" else "QNNExecutionProvider", + device=target.device, + ) + monkeypatch.setattr( session_pkg, "resolve_device", - lambda target: EPDeviceTarget( - ep=target.ep if target.ep != "auto" else "QNNExecutionProvider", - device=target.device, - ), + resolve_target, ) monkeypatch.setattr( session_pkg.WinMLEPRegistry, @@ -104,6 +111,26 @@ def test_explicit_ep_reaches_build_when_compile_is_none( ) +def test_runtime_non_cgc_target_preserves_build_stages( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from winml.modelkit.models import WinMLAutoModel + + received = _install_stubs(monkeypatch, compile_provider="QNNExecutionProvider") + + with pytest.raises(_StopAfterEpCheckError): + WinMLAutoModel.from_pretrained( + "microsoft/resnet-50", + runtime="winml-runtime", + ) + + build_config = received["config"] + assert received["resolve_backend"] == "cgc" + assert build_config.skip_optimize is False + assert build_config.quant is not None + assert build_config.compile is not None + + def test_compile_provider_used_when_user_ep_absent( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -227,11 +254,18 @@ def from_pretrained(*_args: Any, **kwargs: Any) -> str: ) result = WinMLAutoModel.from_pretrained( - "some/composite", task="faketask", allow_unsupported_nodes=True + "some/composite", + task="faketask", + ep_device=MagicMock(), + runtime="winml-runtime", + backend="ort", + allow_unsupported_nodes=True, ) assert result == "COMPOSITE_SENTINEL" assert received.get("allow_unsupported_nodes") is True + assert received.get("runtime") == "winml-runtime" + assert received.get("backend") == "ort" def test_cache_reuse_does_not_eagerly_load_hf_weights( diff --git a/tests/unit/optim/pipes/test_pipe_cgir_rewrite.py b/tests/unit/optim/pipes/test_pipe_cgir_rewrite.py new file mode 100644 index 000000000..6d9a578c2 --- /dev/null +++ b/tests/unit/optim/pipes/test_pipe_cgir_rewrite.py @@ -0,0 +1,1143 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- + +from __future__ import annotations + +from importlib import import_module +from typing import TYPE_CHECKING + +import numpy as np +import pytest +from onnx import ( + ModelProto, + NodeProto, + TensorProto, + checker, + helper, + numpy_helper, + version_converter, +) +from onnx.reference import ReferenceEvaluator + +from winml.modelkit.optim import OptimizationError, Optimizer +from winml.modelkit.optim.pipes import ( + PIPES, + CGIRRewritePipe, + CGIRRewritePipeConfig, + RewritePipe, +) +from winml.modelkit.pattern import PatternMatcher + + +if TYPE_CHECKING: + from collections.abc import Callable + + +def test_initializer_shape_without_value_info(): + from winml.modelkit.pattern.cgc.cgc_constant_folding import _ConstantParameters + + weights = np.random.default_rng(42).normal(size=(2, 3)).astype(np.float32) + model = helper.make_model( + helper.make_graph( + [helper.make_node("Shape", ["weights"], ["shape"])], + "initializer_shape", + [], + [helper.make_tensor_value_info("shape", TensorProto.INT64, [weights.ndim])], + [numpy_helper.from_array(weights, "weights")], + ), + opset_imports=[helper.make_opsetid("", 17)], + ir_version=10, + ) + checker.check_model(model) + expected = ReferenceEvaluator(model).run(None, {})[0] + actual = _ConstantParameters(model, static_shapes=True).evaluate("shape", set(), set()) + np.testing.assert_array_equal(actual, expected) + + +@pytest.mark.parametrize("domain", ["", "com.microsoft"]) +@pytest.mark.parametrize("scale_shape", [(), (1,)]) +@pytest.mark.parametrize("zero_mode", ["scalar", "vector", "omitted", "empty"]) +def test_normalize_int32_dq(domain, scale_shape, zero_mode): + import onnxruntime as ort + + from winml.modelkit.optim import WinMLOptimizationConfig + from winml.modelkit.pattern.cgc import normalize_int32_dq + + random = np.random.default_rng(42) + data = random.integers(-1000, 1000, size=(8,), dtype=np.int32) + scale = random.uniform(0.01, 0.1, size=scale_shape).astype(np.float32) + initializers = [numpy_helper.from_array(data, "data"), numpy_helper.from_array(scale, "scale")] + node_inputs = ["data", "scale"] + if zero_mode in ("scalar", "vector"): + zero = np.zeros(() if zero_mode == "scalar" else (1,), dtype=np.int32) + initializers.append(numpy_helper.from_array(zero, "zero")) + node_inputs.append("zero") + elif zero_mode == "empty": + node_inputs.append("") + model = helper.make_model(helper.make_graph( + [helper.make_node("DequantizeLinear", node_inputs, ["result"], domain=domain)], + "dq", [], [helper.make_tensor_value_info("result", TensorProto.FLOAT, data.shape), + helper.make_tensor_value_info("scale", TensorProto.FLOAT, scale.shape)], + initializers, + ), opset_imports=[helper.make_opsetid("", 18)] + ( + [helper.make_opsetid(domain, 1)] if domain else [] + ), ir_version=10) + original = model.SerializeToString() + result = CGIRRewritePipe().process(model, CGIRRewritePipe.build_config(normalize_int32_dq=True)) + checker.check_model(result) + assert model.SerializeToString() == original + assert result.graph.node[0].domain == domain + assert len(result.graph.node[0].input) == 2 + assert normalize_int32_dq(result) is result + assert WinMLOptimizationConfig.for_cgc()["normalize_int32_dq"] + options = ort.SessionOptions() + options.log_severity_level = 3 + reference = ort.InferenceSession(original, options, providers=["CPUExecutionProvider"]) + candidate = ort.InferenceSession( + result.SerializeToString(), options, providers=["CPUExecutionProvider"], + ) + for actual, expected in zip(candidate.run(None, {}), reference.run(None, {}), strict=True): + np.testing.assert_array_equal(actual, expected) + + +@pytest.mark.parametrize( + "guard", ["scale_input", "zero_input", "nonzero", "dtype", "domain", "per_axis"], +) +def test_normalize_int32_dq_preserves_unsupported(guard): + from winml.modelkit.pattern.cgc import normalize_int32_dq + + random = np.random.default_rng(42) + data = random.integers(1, 10, size=8, dtype=np.int32 if guard != "dtype" else np.uint8) + scale = random.uniform(0.01, 0.1, size=8 if guard == "per_axis" else 1).astype(np.float32) + zero = np.zeros((), dtype=data.dtype) + if guard == "nonzero": + zero = random.integers(1, 10, size=(), dtype=data.dtype) + values = {"data": data, "scale": scale, "zero": zero} + inputs = [helper.make_tensor_value_info(name, numpy_helper.from_array(values[name]).data_type, + values[name].shape) + for name in values if guard == name + "_input"] + model = helper.make_model(helper.make_graph( + [helper.make_node("DequantizeLinear", list(values), ["result"], + domain="custom" if guard == "domain" else "")], + "guard", inputs, [helper.make_tensor_value_info("result", TensorProto.FLOAT, data.shape)], + [numpy_helper.from_array(value, name) for name, value in values.items()], + ), opset_imports=[helper.make_opsetid("", 18)], ir_version=10) + assert normalize_int32_dq(model) is model + + +@pytest.mark.parametrize("opset", [11, 17, 18]) +@pytest.mark.parametrize("shared", [False, True]) +def test_fold_constant_pad_pads(opset, shared): + from winml.modelkit.pattern.cgc import fold_constant_pad_pads + + seed = np.random.default_rng(42) + widths = seed.integers(0, 3, size=4, dtype=np.int32) + outputs = [helper.make_tensor_value_info("result", TensorProto.FLOAT, [None, None])] + if shared: + outputs.append(helper.make_tensor_value_info("pads", TensorProto.INT64, [4])) + model = helper.make_model( + helper.make_graph( + [helper.make_node("Cast", ["widths"], ["pads"], to=TensorProto.INT64), + helper.make_node("Pad", ["source", "pads"], ["result"])], + "constant_pad", [helper.make_tensor_value_info("source", TensorProto.FLOAT, [2, 3])], + outputs, [numpy_helper.from_array(widths, "widths")], + ), + opset_imports=[helper.make_opsetid("", opset)], ir_version=10, + ) + original = model.SerializeToString() + feeds = {"source": seed.normal(size=(2, 3)).astype(np.float32)} + expected = ReferenceEvaluator(model).run(None, feeds) + result = CGIRRewritePipe().process( + model, CGIRRewritePipe.build_config(fold_constant_pad_pads=True), + ) + checker.check_model(result) + assert model.SerializeToString() == original + assert list(result.graph.input) == list(model.graph.input) + assert list(result.graph.output) == list(model.graph.output) + assert sum(node.op_type == "Cast" for node in result.graph.node) == int(shared) + assert fold_constant_pad_pads(result) is result + for actual, reference in zip( + ReferenceEvaluator(result).run(None, feeds), expected, strict=True, + ): + np.testing.assert_array_equal(actual, reference) + + +@pytest.mark.parametrize("opset", [16, 20]) +@pytest.mark.parametrize("align", [0, 1]) +@pytest.mark.parametrize("fp16", [False, True]) +@pytest.mark.parametrize("dynamic", [False, True]) +@pytest.mark.parametrize("spatial,output_hw", [ + ((4, 7), (5, 6)), ((1, 5), (3, 1)), ((5, 1), (1, 4)), ((1, 1), (2, 3)), +]) +def test_gridsample_to_gather(opset, align, fp16, dynamic, spatial, output_hw): + import onnxruntime as ort + + dtype = TensorProto.FLOAT16 if fp16 else TensorProto.FLOAT + batch_dim = "batch" if dynamic else 2 + model = helper.make_model(helper.make_graph([ + helper.make_node("GridSample", ["X", "grid"], ["Y"], + mode="linear" if opset >= 20 else "bilinear", + padding_mode="zeros", align_corners=align), + ], "sampling", [ + helper.make_tensor_value_info("X", dtype, [batch_dim, 3, *spatial]), + helper.make_tensor_value_info("grid", dtype, [batch_dim, *output_hw, 2]), + ], [helper.make_tensor_value_info("Y", dtype, [batch_dim, 3, *output_hw])]), + opset_imports=[helper.make_opsetid("", opset)], ir_version=10) + original = model.SerializeToString() + result = Optimizer().optimize(model, gridsample_to_gather=True) + checker.check_model(result, full_check=True) + assert model.SerializeToString() == original + assert not any(node.op_type == "GridSample" for node in result.graph.node) + assert sum(node.op_type == "GatherND" for node in result.graph.node) == 4 + ranks = {value.name: len(value.type.tensor_type.shape.dim) + for value in result.graph.value_info} + if dynamic: + assert sum(node.op_type == "Range" for node in result.graph.node) == 1 + for node in result.graph.node: + if node.op_type == "GatherND": + assert {attr.name: helper.get_attribute_value(attr) + for attr in node.attribute}["batch_dims"] == 0 + assert all(ranks[name] == 3 for name in [*node.input, *node.output]) + rng = np.random.default_rng(42) + options = ort.SessionOptions() + options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_DISABLE_ALL + reference = ort.InferenceSession(original, options, providers=["CPUExecutionProvider"]) + candidate = ort.InferenceSession( + result.SerializeToString(), options, providers=["CPUExecutionProvider"], + ) + for batch in ([1, 2, 3] if dynamic else [2]): + feeds = {"X": rng.normal(size=(batch, 3, *spatial)).astype( + np.float16 if fp16 else np.float32), + "grid": rng.uniform(-3, 3, size=(batch, *output_hw, 2)).astype( + np.float16 if fp16 else np.float32)} + feeds["grid"].reshape(-1, 2)[:3] = np.linspace(-1, 1, 3)[:, None] + expected = reference.run(None, feeds) + actual = candidate.run(None, feeds) + np.testing.assert_allclose(actual[0], expected[0], atol=2e-3 if fp16 else 2e-6, rtol=2e-3) + + +@pytest.mark.parametrize("mode,padding", [("nearest", "zeros"), ("linear", "border"), + ("cubic", "reflection")]) +def test_gridsample_to_gather_preserves_unsupported_modes(mode, padding): + model = helper.make_model(helper.make_graph([ + helper.make_node("GridSample", ["X", "grid"], ["Y"], mode=mode, padding_mode=padding), + ], "sampling", [ + helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 3, 4, 7]), + helper.make_tensor_value_info("grid", TensorProto.FLOAT, [1, 5, 6, 2]), + ], [helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 3, 5, 6])]), + opset_imports=[helper.make_opsetid("", 20)], ir_version=10) + result = CGIRRewritePipe().process( + model, CGIRRewritePipe.build_config(gridsample_to_gather=True), + ) + assert result is model + + +def _constant_pad_chain(*, axes=False): + rank = 2 + nodes = [] + constants = { + "count": np.asarray([rank], dtype=np.int64), + "widths": np.random.default_rng(7).integers(0, 3, rank, dtype=np.int64), + "matrix": np.asarray([-1, 2], dtype=np.int64), + "start": np.asarray([-1], dtype=np.int64), + "end": np.asarray([np.iinfo(np.int64).min], dtype=np.int64), + "axis": np.asarray([0], dtype=np.int64), + "step": np.asarray([-1], dtype=np.int64), + "vector": np.asarray([-1], dtype=np.int64), + } + for name, value in constants.items(): + nodes.append(helper.make_node("Constant", [], [name], value=numpy_helper.from_array(value))) + nodes.extend([ + helper.make_node( + "ConstantOfShape", ["count"], ["zeros"], + value=numpy_helper.from_array(np.zeros(1, dtype=np.int64)), + ), + helper.make_node("Concat", ["widths", "zeros"], ["joined"], axis=0), + helper.make_node("Reshape", ["joined", "matrix"], ["pairs"]), + helper.make_node("Slice", ["pairs", "start", "end", "axis", "step"], ["reversed"]), + helper.make_node("Transpose", ["reversed"], ["transposed"], perm=[1, 0]), + helper.make_node("Reshape", ["transposed", "vector"], ["flattened"]), + helper.make_node("Cast", ["flattened"], ["pads"], to=TensorProto.INT64), + ]) + pad_inputs = ["source", "pads"] + if axes: + nodes.append(helper.make_node( + "Constant", [], ["pad_axes"], + value=numpy_helper.from_array(np.arange(rank, dtype=np.int64)), + )) + pad_inputs.extend(["", "pad_axes"]) + nodes.append(helper.make_node("Pad", pad_inputs, ["result"])) + return helper.make_model( + helper.make_graph( + nodes, "pad_chain", + [helper.make_tensor_value_info("source", TensorProto.FLOAT, [2, 3])], + [helper.make_tensor_value_info("result", TensorProto.FLOAT, [None, None])], + ), opset_imports=[helper.make_opsetid("", 18 if axes else 17)], ir_version=10, + ) + + +@pytest.mark.parametrize("axes", [False, True]) +@pytest.mark.parametrize("protected", ["none", "capture", "annotation"]) +def test_pad_constant_chain_preserves_values_and_references(axes, protected): + from winml.modelkit.pattern.cgc import fold_constant_pad_pads + + model = _constant_pad_chain(axes=axes) + if protected == "capture": + model.graph.input.append(helper.make_tensor_value_info("condition", TensorProto.BOOL, [])) + branch = helper.make_graph( + [helper.make_node("Identity", ["pads"], ["captured"])], "capture", [], + [helper.make_tensor_value_info("captured", TensorProto.INT64, [4])], + ) + model.graph.node.append(helper.make_node( + "If", ["condition"], ["observed"], then_branch=branch, else_branch=branch, + )) + model.graph.output.append(helper.make_tensor_value_info("observed", TensorProto.INT64, [4])) + if protected == "annotation": + annotation = model.graph.quantization_annotation.add(tensor_name="source") + annotation.quant_parameter_tensor_names.add(key="SCALE_TENSOR", value="pads") + feeds = {"source": np.random.default_rng(21).normal(size=(2, 3)).astype(np.float32)} + if protected == "capture": + feeds["condition"] = np.asarray(True) + expected = ReferenceEvaluator(model).run(None, feeds) + original = model.SerializeToString() + result = fold_constant_pad_pads(model) + checker.check_model(result) + assert result is not model + assert original == model.SerializeToString() + assert any("pads" in node.output for node in result.graph.node) == (protected != "none") + assert fold_constant_pad_pads(result) is result + for actual, reference in zip( + ReferenceEvaluator(result).run(None, feeds), expected, strict=True, + ): + np.testing.assert_array_equal(actual, reference) + + +@pytest.mark.parametrize("reason", [ + "runtime", "overridable", "unsupported", "float", "length", "domain", + "allocation", "nodes", "cache", "invalid_axes", "direct", +]) +def test_pad_constant_folding_rejects_unsafe_candidates(reason, monkeypatch): + from winml.modelkit.pattern.cgc import fold_constant_pad_pads + + folding_module = import_module("winml.modelkit.pattern.cgc.cgc_constant_folding") + + model = _constant_pad_chain(axes=reason == "invalid_axes") + producers = {name: node for node in model.graph.node for name in node.output} + if reason in {"runtime", "overridable"}: + model.graph.node.remove(producers["widths"]) + model.graph.input.append(helper.make_tensor_value_info("widths", TensorProto.INT64, [2])) + if reason == "overridable": + model.graph.initializer.append(numpy_helper.from_array(np.zeros(2, np.int64), "widths")) + elif reason == "unsupported": + producers["pads"].CopyFrom(helper.make_node("Identity", ["flattened"], ["pads"])) + elif reason == "float": + producers["pads"].attribute[0].i = TensorProto.FLOAT + elif reason == "length": + model.graph.input[0].type.tensor_type.shape.dim.add(dim_value=2) + elif reason == "domain": + model.graph.node[-1].domain = "custom" + elif reason == "allocation": + producers["count"].attribute[0].t.CopyFrom( + numpy_helper.from_array(np.asarray([folding_module._MAX_ELEMENTS + 1], np.int64)), + ) + elif reason == "nodes": + monkeypatch.setattr(folding_module, "_MAX_NODES", 2) + elif reason == "cache": + monkeypatch.setattr(folding_module, "_MAX_CACHED_ELEMENTS", 1) + elif reason == "invalid_axes": + producers["pad_axes"].attribute[0].t.CopyFrom( + numpy_helper.from_array(np.zeros(2, np.int64)), + ) + elif reason == "direct": + producers["pads"].CopyFrom(helper.make_node( + "Constant", [], ["pads"], value=numpy_helper.from_array(np.zeros(4, np.int64)), + )) + original = model.SerializeToString() + assert fold_constant_pad_pads(model) is model + assert original == model.SerializeToString() + + +def test_pad_folding_is_enabled_only_by_cgc_defaults(): + from winml.modelkit.optim import WinMLOptimizationConfig + + assert not CGIRRewritePipe.build_config().rules + assert WinMLOptimizationConfig.for_cgc()["cgc_constant_folding"] is True + canonical = CGIRRewritePipe.build_config(cgc_constant_folding=True) + assert canonical == CGIRRewritePipe.build_config(fold_constant_pad_pads=True) + assert canonical == CGIRRewritePipe.build_config( + cgc_constant_folding=True, fold_constant_pad_pads=True, + ) + + +@pytest.mark.parametrize("dynamic", [False, True]) +def test_cgc_constant_folding_static_shape_chain(dynamic): + from winml.modelkit.pattern.cgc import cgc_constant_folding + + shape = ["batch" if dynamic else 2, 3] + model = helper.make_model(helper.make_graph([ + helper.make_node("Shape", ["source"], ["shape"]), + helper.make_node("Gather", ["shape", "axis"], ["batch"], axis=0), + helper.make_node("Unsqueeze", ["batch", "axes"], ["batch_vector"]), + helper.make_node("Concat", ["batch_vector", "tail"], ["target"], axis=0), + helper.make_node("Reshape", ["source", "target"], ["reshaped"]), + helper.make_node("Cast", ["reshaped"], ["result"], to=TensorProto.FLOAT), + ], "shape_chain", [helper.make_tensor_value_info("source", TensorProto.FLOAT16, shape)], + [helper.make_tensor_value_info("result", TensorProto.FLOAT, shape)], [ + numpy_helper.from_array(np.asarray(0, np.int64), "axis"), + numpy_helper.from_array(np.asarray([0], np.int64), "axes"), + numpy_helper.from_array(np.asarray([-1], np.int64), "tail"), + ]), opset_imports=[helper.make_opsetid("", 17)], ir_version=10) + original = model.SerializeToString() + result = CGIRRewritePipe().process( + model, CGIRRewritePipe.build_config(cgc_constant_folding=True), + ) + assert original == model.SerializeToString() + assert any(node.op_type == "Shape" for node in result.graph.node) == dynamic + assert sum(node.op_type == "Cast" for node in result.graph.node) == 1 + assert (result is model) == dynamic + assert cgc_constant_folding(result) is result + for batch in ([1, 4] if dynamic else [2]): + feeds = {"source": np.random.default_rng(42).normal(size=(batch, 3)).astype(np.float16)} + np.testing.assert_array_equal( + ReferenceEvaluator(result).run(None, feeds)[0], + ReferenceEvaluator(model).run(None, feeds)[0], + ) + + +@pytest.mark.parametrize("start,end", [(1, 3), (-2, 100), (2, 1)]) +def test_cgc_constant_folding_partial_shape(start, end): + from winml.modelkit.pattern.cgc import cgc_constant_folding + + shape = ["batch", 3, 4] + output_rank = len(shape[start:end]) + model = helper.make_model(helper.make_graph([ + helper.make_node("Shape", ["source"], ["result"], start=start, end=end), + ], "partial_shape", [helper.make_tensor_value_info("source", TensorProto.FLOAT, shape)], + [helper.make_tensor_value_info("result", TensorProto.INT64, [output_rank])]), + opset_imports=[helper.make_opsetid("", 15)], ir_version=10) + result = cgc_constant_folding(model) + assert result.graph.node[0].op_type == "Constant" + for batch in [1, 2]: + feeds = {"source": np.random.default_rng(42).normal(size=(batch, 3, 4)).astype(np.float32)} + np.testing.assert_array_equal( + ReferenceEvaluator(result).run(None, feeds)[0], + ReferenceEvaluator(model).run(None, feeds)[0], + ) + + +def test_cgc_constant_folding_rejects_broadcast_allocation(monkeypatch): + from winml.modelkit.pattern.cgc.cgc_constant_folding import _ConstantParameters + + folding_module = import_module("winml.modelkit.pattern.cgc.cgc_constant_folding") + monkeypatch.setattr(folding_module, "_MAX_ELEMENTS", 32) + model = helper.make_model(helper.make_graph([ + helper.make_node("Add", ["rows", "columns"], ["result"]), + ], "broadcast", [], [helper.make_tensor_value_info("result", TensorProto.INT64, [8, 8])], [ + numpy_helper.from_array(np.arange(8, dtype=np.int64).reshape(8, 1), "rows"), + numpy_helper.from_array(np.arange(8, dtype=np.int64).reshape(1, 8), "columns"), + ]), opset_imports=[helper.make_opsetid("", 17)], ir_version=10) + with pytest.raises(ValueError, match="Broadcast allocation budget"): + _ConstantParameters(model, static_shapes=True).evaluate("result", set(), set()) + + +@pytest.mark.parametrize("source_op", ["Relu", "PRelu"]) +@pytest.mark.parametrize("model_barrier", [False, True]) +def test_cgir_reuses_matcher_without_skipping_rules(source_op, model_barrier, monkeypatch): + import winml.modelkit.optim.pipes.cgir_rewrite as cgir_module + + prepared = [] + matched = [] + + class TrackingMatcher(PatternMatcher): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + prepared.append(self) + + def match(self): + matched.append(tuple(self.patterns)) + return super().match() + + monkeypatch.setattr(cgir_module, "PatternMatcher", TrackingMatcher) + model = helper.make_model( + helper.make_graph( + [helper.make_node( + source_op, ["source", "slope"] if source_op == "PRelu" else ["source"], + ["result"], + )], + "matcher_reuse", + [helper.make_tensor_value_info("source", TensorProto.FLOAT, [2, 3])], + [helper.make_tensor_value_info("result", TensorProto.FLOAT, [2, 3])], + [numpy_helper.from_array( + np.random.default_rng(42).uniform(size=()).astype(np.float32), "slope", + )], + ), + opset_imports=[helper.make_opsetid("", 18)], + ir_version=10, + ) + resize_rule = CGIRRewritePipe.build_config(omit_empty_resize_inputs=True).rules[0] + prelu_rule = CGIRRewritePipe.build_config(prelu_to_relu=True).rules[0] + rules = [resize_rule] + if model_barrier: + rules.extend(CGIRRewritePipe.build_config(eliminate_identity=True).rules) + rules.extend([prelu_rule, resize_rule]) + original = model.SerializeToString() + feeds = {"source": np.exp(np.random.default_rng(42).normal(size=(2, 3))).astype(np.float32)} + expected = ReferenceEvaluator(model).run(None, feeds) + + result = CGIRRewritePipe().process(model, CGIRRewritePipeConfig(rules=rules)) + + assert len(prepared) == 1 + int(model_barrier) + int(source_op == "PRelu") + expected_patterns = [resize_rule.source.__name__, prelu_rule.source.__name__] + if source_op == "PRelu": + expected_patterns.append(prelu_rule.source.__name__) + expected_patterns.append(resize_rule.source.__name__) + assert matched == [(name,) for name in expected_patterns] + assert model.SerializeToString() == original + checker.check_model(result) + for actual, reference in zip( + ReferenceEvaluator(result).run(None, feeds), expected, strict=True, + ): + np.testing.assert_allclose(actual, reference) + + +@pytest.mark.parametrize("capability", [ + "log-to-reduce-log-sum", + "materialize-initializer-parameters", + "fold-scalar-initializer-casts", +]) +def test_retired_cgir_rules_are_not_registered(capability): + from click.testing import CliRunner + + from winml.modelkit.commands.optimize import optimize + from winml.modelkit.optim import WinMLOptimizationConfig, get_all_capabilities + + assert capability not in get_all_capabilities() + assert capability not in CGIRRewritePipe.capabilities + assert capability.replace("-", "_") not in WinMLOptimizationConfig.for_cgc() + result = CliRunner().invoke(optimize, ["--help"]) + assert result.exit_code == 0 + assert f"--enable-{capability}" not in result.output + assert f"--disable-{capability}" not in result.output + + +@pytest.mark.parametrize("enabled", [False, True]) +@pytest.mark.parametrize("protected_output", [False, True]) +def test_eliminate_identity_preserves_results_and_model( + enabled: bool, protected_output: bool, +) -> None: + shape = [2, 3] + nodes = [ + helper.make_node("Identity", ["source"], ["alias"]), + helper.make_node("Identity", ["alias"], ["second_alias"]), + helper.make_node("Add", ["second_alias", "source"], ["result"]), + ] + outputs = [helper.make_tensor_value_info("result", TensorProto.FLOAT, shape)] + if protected_output: + outputs.append(helper.make_tensor_value_info("alias", TensorProto.FLOAT, shape)) + model = helper.make_model( + helper.make_graph( + nodes, "aliases", + [helper.make_tensor_value_info("source", TensorProto.FLOAT, shape)], + outputs, + value_info=[ + helper.make_tensor_value_info("alias", TensorProto.FLOAT, shape), + helper.make_tensor_value_info("second_alias", TensorProto.FLOAT, shape), + ], + ), + opset_imports=[helper.make_opsetid("", 18)], + ir_version=11, + ) + original = model.SerializeToString() + feeds = {"source": np.random.default_rng(42).standard_normal(shape).astype(np.float32)} + expected = ReferenceEvaluator(model).run(None, feeds) + + result = CGIRRewritePipe().process( + model, CGIRRewritePipe.build_config(eliminate_identity=enabled), + ) + + checker.check_model(result) + assert model.SerializeToString() == original + assert result.graph.input == model.graph.input + assert result.graph.output == model.graph.output + retained = [node for node in result.graph.node if node.op_type == "Identity"] + assert len(retained) == (0 if enabled else 2) + assert sum(node.op_type == "Reshape" for node in result.graph.node) == int( + enabled and protected_output, + ) + for actual, reference in zip( + ReferenceEvaluator(result).run(None, feeds), expected, strict=True, + ): + np.testing.assert_array_equal(actual, reference) + repeated = CGIRRewritePipe().process( + result, CGIRRewritePipe.build_config(eliminate_identity=enabled), + ) + assert repeated.SerializeToString() == result.SerializeToString() + + +@pytest.mark.parametrize("guard", ["none", "dynamic", "zero", "scalar", "fp16", "annotation", + "subgraph", "mismatch", "unknown", "opset"]) +def test_identity_output_scope_and_bits(guard): + from winml.modelkit.pattern.cgc import eliminate_identity + + shape = {"dynamic": ["batch"], "zero": [0], "scalar": []}.get(guard, [8]) + dtype = TensorProto.FLOAT16 if guard == "fp16" else TensorProto.FLOAT + model = helper.make_model(helper.make_graph( + [helper.make_node("Identity", ["source"], ["result"])], "output_alias", + [helper.make_tensor_value_info("source", dtype, shape)], + [helper.make_tensor_value_info("result", dtype, shape)], + ), opset_imports=[helper.make_opsetid("", 4 if guard == "opset" else 18)], ir_version=10) + if guard == "annotation": + model.graph.quantization_annotation.add(tensor_name="result") + elif guard == "subgraph": + branch = helper.make_graph([], "branch", [], []) + model.graph.node.append(helper.make_node("If", ["cond"], ["other"], then_branch=branch)) + elif guard == "mismatch": + model.graph.value_info.append(helper.make_tensor_value_info("source", dtype, [9])) + elif guard == "unknown": + model.graph.input[0].type.tensor_type.ClearField("shape") + original = model.SerializeToString() + result = eliminate_identity(model) + assert model.SerializeToString() == original + if guard != "none": + assert result is model + return + checker.check_model(result, full_check=True) + assert result.graph.output == model.graph.output + assert result.graph.input == model.graph.input + assert result.graph.node[0].op_type == "Reshape" + assert eliminate_identity(result) is result + random = np.random.default_rng(42) + values = [random.normal(size=shape).astype(np.float32)] + values.extend(np.full(shape, special, dtype=np.float32) for special in ( + 0.0, -0.0, np.inf, -np.inf, np.nan, + np.nextafter(np.float32(0), np.float32(1)), + )) + for data in values: + expected = ReferenceEvaluator(model).run(None, {"source": data})[0] + actual = ReferenceEvaluator(result).run(None, {"source": data})[0] + assert actual.tobytes() == expected.tobytes() + + + + + + + + + + + + +def _make_resize_model( + *, + opset: int, + roi: np.ndarray, + scales: np.ndarray, + use_sizes: bool, + cast_inputs: bool = True, + resize_count: int = 1, +) -> ModelProto: + nodes: list[NodeProto] = [] + initializers = [ + numpy_helper.from_array(roi, "roi_source"), + numpy_helper.from_array(scales, "scales_source"), + ] + roi_name = "roi_source" + scales_name = "scales_source" + if cast_inputs: + nodes.extend( + [ + helper.make_node("Cast", [roi_name], ["roi"], to=TensorProto.FLOAT), + helper.make_node( + "Cast", + [scales_name], + ["scales"], + to=TensorProto.FLOAT, + ), + ] + ) + roi_name = "roi" + scales_name = "scales" + + outputs = [] + for index in range(resize_count): + node_inputs = ["X", roi_name, scales_name] + if use_sizes: + sizes_name = f"sizes_{index}" + initializers.append( + numpy_helper.from_array( + np.array([1, 1, 4, 4], dtype=np.int64), + sizes_name, + ) + ) + node_inputs.append(sizes_name) + output_name = f"Y_{index}" + nodes.append( + helper.make_node( + "Resize", + node_inputs, + [output_name], + name=f"resize_{index}", + mode="nearest", + ) + ) + outputs.append( + helper.make_tensor_value_info( + output_name, + TensorProto.FLOAT, + [1, 1, 4, 4], + ) + ) + + graph = helper.make_graph( + nodes, + "resize", + [helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 1, 2, 2])], + outputs, + initializers, + ) + return helper.make_model( + graph, + opset_imports=[helper.make_opsetid("", opset)], + ir_version=11, + ) + + +def _make_tile_model( + repeats: np.ndarray, + *, + initializer_backed: bool, +) -> ModelProto: + input_shape = [2, 3] + output_shape = [ + dimension * int(repeat) + for dimension, repeat in zip(input_shape, repeats, strict=True) + ] + repeats_tensor = numpy_helper.from_array(repeats, "repeats") + nodes = [] + initializers = [] + if initializer_backed: + initializers.append(repeats_tensor) + else: + nodes.append( + helper.make_node( + "Constant", + [], + ["repeats"], + name="repeats_constant", + value=repeats_tensor, + ) + ) + nodes.append(helper.make_node("Tile", ["X", "repeats"], ["Y"], name="tile")) + return helper.make_model( + helper.make_graph( + nodes, + "tile", + [helper.make_tensor_value_info("X", TensorProto.FLOAT, input_shape)], + [helper.make_tensor_value_info("Y", TensorProto.FLOAT, output_shape)], + initializer=initializers, + ), + opset_imports=[helper.make_opsetid("", 18)], + ir_version=11, + ) + + + + + + +def _empty_float() -> np.ndarray: + return np.empty((0,), dtype=np.float32) + + +def _nonempty_roi() -> np.ndarray: + return np.array([0.0, 0.0, 1.0, 1.0], dtype=np.float32) + + +def _effective_scales() -> np.ndarray: + return np.array([1.0, 1.0, 2.0, 2.0], dtype=np.float32) + + +def _enabled_config() -> CGIRRewritePipeConfig: + return CGIRRewritePipe.build_config(omit_empty_resize_inputs=True) + + +def _resize_nodes(model: ModelProto) -> list[NodeProto]: + return [node for node in model.graph.node if node.op_type == "Resize"] + + +def _default_opset(model: ModelProto) -> int: + return next( + int(opset.version) + for opset in model.opset_import + if opset.domain in {"", "ai.onnx"} + ) + + +def test_omit_empty_resize_inputs_is_disabled_by_default() -> None: + config = CGIRRewritePipe.build_config() + + assert config.rules == [] + assert not CGIRRewritePipe.should_process(config) + assert CGIRRewritePipe.capabilities["omit-empty-resize-inputs"].default is False + + + + +def test_disable_graph_optimization_does_not_enable_rules_implicitly() -> None: + config = CGIRRewritePipe.build_config(ort_graph_optimization=False) + + assert config.rules == [] + assert CGIRRewritePipe.should_process(config) is False + + +def test_omit_empty_resize_inputs_is_owned_only_by_cgir_rewrite_pipe() -> None: + assert "omit-empty-resize-inputs" in CGIRRewritePipe.capabilities + assert "omit-empty-resize-inputs" not in RewritePipe.capabilities + assert PIPES[0] is CGIRRewritePipe + + + + + + +def test_omit_empty_resize_inputs_requires_explicit_enable() -> None: + config = _enabled_config() + + assert CGIRRewritePipe.should_process(config) + assert len(config.rules) == 1 + + +@pytest.mark.parametrize("opset", [11, 12]) +def test_omit_empty_resize_inputs_updates_legacy_opset_and_rewrites(opset: int) -> None: + model = _make_resize_model( + opset=opset, + roi=_empty_float(), + scales=_empty_float(), + use_sizes=True, + ) + + result = CGIRRewritePipe().process(model, _enabled_config()) + + assert _default_opset(model) == opset + assert list(_resize_nodes(model)[0].input) == ["X", "roi", "scales", "sizes_0"] + assert _default_opset(result) == 13 + assert list(_resize_nodes(result)[0].input) == ["X", "", "", "sizes_0"] + checker.check_model(result) + + +def test_omit_empty_resize_inputs_rewrites_opset_13_without_conversion( + monkeypatch: pytest.MonkeyPatch, +) -> None: + model = _make_resize_model( + opset=13, + roi=_empty_float(), + scales=_empty_float(), + use_sizes=True, + ) + + def unexpected_conversion(*_args: object, **_kwargs: object) -> ModelProto: + pytest.fail("opset 13 model must not be converted") + + monkeypatch.setattr( + "winml.modelkit.optim.pipes.cgir_rewrite.version_converter.convert_version", + unexpected_conversion, + ) + + result = CGIRRewritePipe().process(model, _enabled_config()) + + assert _default_opset(result) == 13 + assert list(_resize_nodes(result)[0].input) == ["X", "", "", "sizes_0"] + + +def test_omit_empty_resize_inputs_rewrites_all_matches() -> None: + model = _make_resize_model( + opset=11, + roi=_empty_float(), + scales=_empty_float(), + use_sizes=True, + resize_count=3, + ) + + result = CGIRRewritePipe().process(model, _enabled_config()) + + assert [ + list(node.input) + for node in _resize_nodes(result) + ] == [ + ["X", "", "", "sizes_0"], + ["X", "", "", "sizes_1"], + ["X", "", "", "sizes_2"], + ] + + +def test_omit_empty_resize_inputs_preserves_effective_scales() -> None: + model = _make_resize_model( + opset=13, + roi=_empty_float(), + scales=_effective_scales(), + use_sizes=False, + ) + + result = CGIRRewritePipe().process(model, _enabled_config()) + + assert list(_resize_nodes(result)[0].input) == ["X", "", "scales"] + + +@pytest.mark.parametrize( + "model_factory", + [ + lambda: _make_resize_model( + opset=11, + roi=_nonempty_roi(), + scales=_empty_float(), + use_sizes=True, + ), + lambda: _make_resize_model( + opset=11, + roi=_nonempty_roi(), + scales=_effective_scales(), + use_sizes=True, + ), + lambda: _make_resize_model( + opset=13, + roi=_nonempty_roi(), + scales=_effective_scales(), + use_sizes=False, + ), + ], + ids=["nonempty-roi", "scales-and-sizes", "no-empty-inputs"], +) +def test_omit_empty_resize_inputs_does_not_rewrite_nonmatches( + model_factory: Callable[[], ModelProto], + monkeypatch: pytest.MonkeyPatch, +) -> None: + model = model_factory() + original = model.SerializeToString() + + def unexpected_conversion(*_args: object, **_kwargs: object) -> ModelProto: + pytest.fail("non-matching model must not be converted") + + monkeypatch.setattr( + "winml.modelkit.optim.pipes.cgir_rewrite.version_converter.convert_version", + unexpected_conversion, + ) + + result = CGIRRewritePipe().process(model, _enabled_config()) + + assert result is model + assert result.SerializeToString() == original + + +def test_omit_empty_resize_inputs_does_not_rewrite_omitted_inputs() -> None: + model = _make_resize_model( + opset=13, + roi=_empty_float(), + scales=_empty_float(), + use_sizes=True, + ) + resize = _resize_nodes(model)[0] + resize.input[1] = "" + resize.input[2] = "" + original = model.SerializeToString() + + result = CGIRRewritePipe().process(model, _enabled_config()) + + assert result is model + assert result.SerializeToString() == original + + +def test_omit_empty_resize_inputs_converts_once( + monkeypatch: pytest.MonkeyPatch, +) -> None: + model = _make_resize_model( + opset=11, + roi=_empty_float(), + scales=_empty_float(), + use_sizes=True, + resize_count=2, + ) + calls: list[int] = [] + convert_version = version_converter.convert_version + + def recorded_conversion( + source: ModelProto, + target_opset: int, + ) -> ModelProto: + calls.append(target_opset) + return convert_version(source, target_opset) + + monkeypatch.setattr( + "winml.modelkit.optim.pipes.cgir_rewrite.version_converter.convert_version", + recorded_conversion, + ) + + CGIRRewritePipe().process(model, _enabled_config()) + + assert calls == [13] + + +def test_omit_empty_resize_inputs_surfaces_conversion_failure( + monkeypatch: pytest.MonkeyPatch, +) -> None: + model = _make_resize_model( + opset=11, + roi=_empty_float(), + scales=_empty_float(), + use_sizes=True, + ) + original = model.SerializeToString() + + def failed_conversion(*_args: object, **_kwargs: object) -> ModelProto: + raise RuntimeError("conversion failed") + + monkeypatch.setattr( + "winml.modelkit.optim.pipes.cgir_rewrite.version_converter.convert_version", + failed_conversion, + ) + + with pytest.raises(OptimizationError, match="conversion failed"): + CGIRRewritePipe().process(model, _enabled_config()) + + assert model.SerializeToString() == original + + +def test_optimizer_leaves_resize_unchanged_by_default() -> None: + model = _make_resize_model( + opset=11, + roi=_empty_float(), + scales=_empty_float(), + use_sizes=True, + ) + + result = Optimizer().optimize(model) + + assert _default_opset(result) == 11 + resize_inputs = list(_resize_nodes(result)[0].input) + assert resize_inputs[0] == "X" + assert resize_inputs[1] + assert resize_inputs[2] + assert resize_inputs[3] == "sizes_0" + + +def test_optimizer_applies_explicit_cgir_resize_rewrite() -> None: + model = _make_resize_model( + opset=11, + roi=_empty_float(), + scales=_empty_float(), + use_sizes=True, + ) + + result = Optimizer().optimize(model, omit_empty_resize_inputs=True) + + assert _default_opset(result) == 13 + assert list(_resize_nodes(result)[0].input) == ["X", "", "", "sizes_0"] + checker.check_model(result) + + + + + + + + + + + + + + + + +def test_deduplicate_opset_imports_is_explicit_and_cgir_only() -> None: + capability = "deduplicate-opset-imports" + assert CGIRRewritePipe.capabilities[capability].default is False + assert capability not in RewritePipe.capabilities + model = _make_tile_model(np.ones(2, dtype=np.int64), initializer_backed=True) + model.opset_import.append(model.opset_import[0]) + config = CGIRRewritePipe.build_config(ort_graph_optimization=False) + assert CGIRRewritePipe().process(model, config) is model + enabled = CGIRRewritePipe.build_config(deduplicate_opset_imports=True) + assert CGIRRewritePipe.should_process(enabled) + assert len(enabled.rules) == 1 + + +def test_deduplicate_opset_imports_preserves_model_content_and_domain_order() -> None: + model = _make_tile_model(np.ones(2, dtype=np.int64), initializer_backed=True) + model.opset_import.extend( + [ + helper.make_opsetid("example.first", 1), + model.opset_import[0], + helper.make_opsetid("example.second", 1), + helper.make_opsetid("example.first", 1), + ] + ) + model.functions.append( + helper.make_function( + "example.first", "PassThrough", ["X"], ["Y"], + [helper.make_node("Identity", ["X"], ["Y"])], + [helper.make_opsetid("", 18)], + ) + ) + original = model.SerializeToString() + expected_imports = list(dict.fromkeys( + (opset.domain, opset.version) for opset in model.opset_import + )) + result = CGIRRewritePipe().process( + model, CGIRRewritePipe.build_config(deduplicate_opset_imports=True) + ) + assert [(opset.domain, opset.version) for opset in result.opset_import] == expected_imports + assert result.graph.SerializeToString() == model.graph.SerializeToString() + without_imports = ModelProto() + without_imports.CopyFrom(result) + del without_imports.opset_import[:] + without_imports.opset_import.extend(model.opset_import) + assert without_imports.SerializeToString() == original + assert model.SerializeToString() == original + checker.check_model(result) + config = CGIRRewritePipe.build_config(deduplicate_opset_imports=True) + assert CGIRRewritePipe().process(result, config) is result + + +@pytest.mark.parametrize("domain", ["", "example.custom"]) +def test_deduplicate_opset_imports_rejects_conflicts_without_mutation(domain: str) -> None: + model = _make_tile_model(np.ones(2, dtype=np.int64), initializer_backed=True) + model.opset_import.extend( + [helper.make_opsetid(domain, 18), helper.make_opsetid(domain, 19)] + ) + original = model.SerializeToString() + with pytest.raises(OptimizationError, match="Conflicting opset imports"): + CGIRRewritePipe().process( + model, CGIRRewritePipe.build_config(deduplicate_opset_imports=True) + ) + assert model.SerializeToString() == original + + +def test_deduplicate_opset_imports_precedes_version_conversion() -> None: + model = _make_resize_model( + opset=11, roi=_empty_float(), scales=_empty_float(), use_sizes=True + ) + model.opset_import.append(model.opset_import[0]) + original = model.SerializeToString() + result = CGIRRewritePipe().process( + model, + CGIRRewritePipe.build_config( + omit_empty_resize_inputs=True, deduplicate_opset_imports=True + ), + ) + assert [opset.version for opset in result.opset_import if opset.domain == ""] == [13] + assert list(_resize_nodes(result)[0].input) == ["X", "", "", "sizes_0"] + assert model.SerializeToString() == original + checker.check_model(result) + + +def test_optimizer_applies_opset_deduplication_without_export() -> None: + model = _make_tile_model(np.ones(2, dtype=np.int64), initializer_backed=True) + model.opset_import.append(model.opset_import[0]) + result = Optimizer().optimize( + model, ort_graph_optimization=False, deduplicate_opset_imports=True + ) + assert len(result.opset_import) == 1 + assert _default_opset(result) == _default_opset(model) diff --git a/tests/unit/optim/pipes/test_pipe_config.py b/tests/unit/optim/pipes/test_pipe_config.py index 8f0be775f..67edc7780 100644 --- a/tests/unit/optim/pipes/test_pipe_config.py +++ b/tests/unit/optim/pipes/test_pipe_config.py @@ -373,6 +373,12 @@ def test_optimization_level_kwarg_ignored(self) -> None: assert config.optimization_level == 2 + def test_disable_graph_optimization_disables_graph_pipe(self) -> None: + config = ORTGraphPipe.build_config(ort_graph_optimization=False, gelu_fusion=True) + + assert config.optimization_level == 0 + assert ORTGraphPipe.should_process(config) is False + def test_verbose_kwarg_passed_through(self) -> None: """verbose kwarg is passed to config.""" config = ORTGraphPipe.build_config(verbose=True) @@ -393,9 +399,12 @@ def test_graph_capabilities_same_as_pipe(self) -> None: assert ORTGraphPipe.capabilities is GRAPH_CAPABILITIES def test_all_caps_have_ort_name(self) -> None: - """All capabilities in GRAPH_CAPABILITIES have ort_name.""" + """Optimizer capabilities have ORT names; the stage switch does not.""" for name, cap in GRAPH_CAPABILITIES.items(): assert hasattr(cap, "ort_name"), f"{name} missing ort_name" + if name == "ort-graph-optimization": + assert cap.ort_name is None + continue assert cap.ort_name, f"{name} has empty ort_name" def test_most_bool_caps_are_default_false(self) -> None: @@ -409,7 +418,7 @@ def test_most_bool_caps_are_default_false(self) -> None: from winml.modelkit.optim.registry import BoolCapability # Capabilities that are allowed to have default=True - allowed_default_true = {"constant-folding"} + allowed_default_true = {"constant-folding", "ort-graph-optimization"} for name, cap in GRAPH_CAPABILITIES.items(): if isinstance(cap, BoolCapability): diff --git a/tests/unit/optim/test_api.py b/tests/unit/optim/test_api.py index 1c2c1d0ff..03ff24f5c 100644 --- a/tests/unit/optim/test_api.py +++ b/tests/unit/optim/test_api.py @@ -102,6 +102,39 @@ def mock_capability() -> MagicMock: # ============================================================================= +def test_cgc_build_config_skips_ort_graph(simple_model): + from winml.modelkit.optim import WinMLOptimizationConfig + from winml.modelkit.optim.pipes import ORTGraphPipe + + config = WinMLOptimizationConfig.from_dict( + WinMLOptimizationConfig.for_cgc().to_dict() + ) + with patch.object(ORTGraphPipe, "process", side_effect=AssertionError("ORT graph ran")): + result = optimize_onnx(simple_model, **config) + onnx.checker.check_model(result) + + +@pytest.mark.parametrize("enabled", [True, False]) +@pytest.mark.parametrize("source", ["kwargs", "config"]) +def test_graph_optimization_forwarded_to_optimizer(simple_model, enabled, source): + with patch("winml.modelkit.optim.api.Optimizer") as optimizer: + optimizer.return_value.optimize.return_value = simple_model + if source == "config": + optimize_onnx(simple_model, config={"ort-graph-optimization": enabled}) + else: + optimize_onnx(simple_model, ort_graph_optimization=enabled) + for call in optimizer.return_value.optimize.call_args_list: + assert call.kwargs["ort_graph_optimization"] is enabled + assert "backend" not in call.kwargs + + +def test_graph_optimization_enabled_by_default(simple_model): + with patch("winml.modelkit.optim.api.Optimizer") as optimizer: + optimizer.return_value.optimize.return_value = simple_model + optimize_onnx(simple_model) + assert optimizer.return_value.optimize.call_args.kwargs["ort_graph_optimization"] is True + + class TestLoadModel: """Tests for _load_model helper function.""" diff --git a/tests/unit/optim/test_optimizer.py b/tests/unit/optim/test_optimizer.py index 3095e30a7..e8b275838 100644 --- a/tests/unit/optim/test_optimizer.py +++ b/tests/unit/optim/test_optimizer.py @@ -699,8 +699,8 @@ def test_registered_pipes_count(self) -> None: """Verify the expected number of pipes are registered.""" Optimizer._initialize_pipes() # Currently: ORTGraphPipe, AlgebraicRewritePipe, RewritePipe, - # ORTFusionPipe, SurgeryPipe - assert len(Optimizer.pipes) == 5 + # ORTFusionPipe, SurgeryPipe, CGIRRewritePipe + assert len(Optimizer.pipes) == 6 def test_registered_pipe_names(self) -> None: """Verify expected pipe names are registered.""" diff --git a/tests/unit/session/test_ep_device.py b/tests/unit/session/test_ep_device.py index 2d088f947..e52e817dc 100644 --- a/tests/unit/session/test_ep_device.py +++ b/tests/unit/session/test_ep_device.py @@ -237,6 +237,31 @@ def test_resolve_device_does_not_load_dll() -> None: mock_reg.instance.assert_not_called() +def test_resolve_device_cgc_loads_runtime_before_resolving_with_dml() -> None: + calls: list[object] = [] + registry = MagicMock() + registry.available_eps.return_value = frozenset({"DmlExecutionProvider"}) + registry.auto_device.side_effect = lambda target: calls.append(target) + + with ( + patch( + "winml.modelkit.session._runtime_import.import_runtime", + side_effect=lambda: calls.append("runtime"), + ), + patch( + "winml.modelkit.session.ep_registry.WinMLEPRegistry.instance", + return_value=registry, + ), + ): + result = resolve_device( + EPDeviceTarget(ep="auto", device="auto", source="pypi"), + backend="cgc", + ) + + assert calls == ["runtime", EPDeviceTarget(ep="DmlExecutionProvider", device="gpu")] + assert result == EPDeviceTarget(ep="DmlExecutionProvider", device="gpu") + + @pytest.mark.parametrize( "ep,device", [ @@ -420,7 +445,7 @@ def test_ep_device_specs_count() -> None: """ from winml.modelkit.session import EP_DEVICE_SPECS - assert len(EP_DEVICE_SPECS) == 12 + assert len(EP_DEVICE_SPECS) == 13 def test_lookup_device_spec_qnn_npu() -> None: diff --git a/tests/unit/session/test_runtime_session.py b/tests/unit/session/test_runtime_session.py new file mode 100644 index 000000000..e93d0c389 --- /dev/null +++ b/tests/unit/session/test_runtime_session.py @@ -0,0 +1,857 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Unit tests for the Windows ML Runtime inference backend.""" + +from __future__ import annotations + +import json +import threading +import time +from types import SimpleNamespace +from typing import TYPE_CHECKING, Any +from unittest.mock import Mock + +import click +import numpy as np +import pytest + +from winml.modelkit.session.runtime_session import ( + WinMLRuntimeSession, + _apply_io_metadata, + _DXCoreAdapter, + _numpy_dtype_for, + _ResolvedRuntimeTarget, + _shape_with_dynamic_dims, +) + + +if TYPE_CHECKING: + from pathlib import Path + + +class _NotSupportedError(Exception): + pass + + +class _Schema: + input_count = 1 + output_count = 1 + + def input_desc(self, index: int) -> tuple[Any, list[int]]: + assert index == 0 + return SimpleNamespace(name="FLOAT32"), [0, 3, 8, 8] + + def output_desc(self, index: int) -> tuple[Any, list[int]]: + assert index == 0 + return SimpleNamespace(name="FLOAT32"), [0, 5] + + +class _Tensor: + def __init__(self, value: np.ndarray) -> None: + self._value = value + + def to_numpy(self) -> np.ndarray: + return self._value + + +class _Stage: + def __init__(self) -> None: + self.execution_target = SimpleNamespace(kind=SimpleNamespace(name="GPU")) + self.bound: dict[int, _Tensor] = {} + self.requested_outputs: set[int] = set() + self.caller_owns_state_tensors = False + + def schema(self) -> _Schema: + return _Schema() + + def ort_diagnostics(self) -> None: + raise _NotSupportedError + + def bind_input(self, index: int, tensor: _Tensor) -> None: + self.bound[index] = tensor + + def output(self, index: int) -> _Tensor: + assert index == 0 + assert index in self.requested_outputs + batch = self.bound[0].to_numpy().shape[0] + return _Tensor(np.zeros((batch, 5), dtype=np.float32)) + + def request_output(self, index: int) -> None: + self.requested_outputs.add(index) + + def close(self) -> None: + pass + + +class _Pipeline: + def __init__(self) -> None: + self.runs = 0 + + def run(self) -> None: + self.runs += 1 + + def close(self) -> None: + pass + + +class _Builder: + def __init__(self, stage: _Stage, pipeline: _Pipeline) -> None: + self.stage = stage + self.pipeline = pipeline + self.targets: list[Any] = [] + self.caller_owned_at_build: bool | None = None + + def add_model_stage(self, model: Any, target: Any = None) -> _Stage: + self.targets.append(target) + return self.stage + + def build(self) -> _Pipeline: + self.caller_owned_at_build = self.stage.caller_owns_state_tensors + return self.pipeline + + +class _Model: + def ort_schema(self) -> None: + raise _NotSupportedError + + def close(self) -> None: + pass + + +class _Runtime: + def __init__(self, stage: _Stage, pipeline: _Pipeline) -> None: + self.builder = _Builder(stage, pipeline) + + def load_model(self, path: str) -> _Model: + assert path.endswith(".mlir") + return _Model() + + def create_pipeline_builder(self) -> _Builder: + return self.builder + + def create_target_from_adapter(self, adapter: object) -> object: + return adapter + + def tensor_from_numpy(self, value: np.ndarray) -> _Tensor: + return _Tensor(value) + + def close(self) -> None: + pass + + +class _AdapterHandle: + pointer = 123 + + def __init__(self) -> None: + self.closed = False + + def close(self) -> None: + self.closed = True + + +def _mlir_ep_device() -> SimpleNamespace: + return SimpleNamespace( + device=SimpleNamespace( + device_type="GPU", + hardware_name="Test GPU", + adapter_luid=456, + ) + ) + + +@pytest.fixture +def mlir_target( + monkeypatch: pytest.MonkeyPatch, +) -> tuple[SimpleNamespace, _AdapterHandle]: + adapter = _AdapterHandle() + + def from_luid(cls: type[_DXCoreAdapter], luid: int) -> _AdapterHandle: + assert luid == 456 + return adapter + + monkeypatch.setattr(_DXCoreAdapter, "from_luid", classmethod(from_luid)) + return _mlir_ep_device(), adapter + + +def test_onnx_cgc_reloads_compiled_model_without_io_counts() -> None: + source_model = Mock() + compiled_model = _Model() + compiler = Mock() + loaded_paths: list[str] = [] + + def load_model(path: str) -> Any: + loaded_paths.append(path) + return source_model if len(loaded_paths) == 1 else compiled_model + + runtime = SimpleNamespace(load_model=load_model) + target = _ResolvedRuntimeTarget( + execution_target=SimpleNamespace(model_compiler=lambda: compiler), + device_class="gpu", + ) + session = WinMLRuntimeSession("model.onnx", ep_device=_mlir_ep_device(), backend="cgc") + try: + model, ort_schema, has_named_bindings = session._load_onnx_on_cgc(runtime, target) + + assert model is compiled_model + assert ort_schema is source_model.ort_schema.return_value + assert has_named_bindings is False + assert loaded_paths[0] == str(session.running_model_path) + assert len(loaded_paths) == 2 + compiler.compile_to_file.assert_called_once_with(source_model, loaded_paths[1]) + compiler.close.assert_called_once_with() + source_model.schema.assert_not_called() + finally: + session.reset() + + +def test_schema_helpers() -> None: + assert _numpy_dtype_for(SimpleNamespace(name="FLOAT16")) == "float16" + assert _shape_with_dynamic_dims([0, -1, (1 << 64) - 1, 4]) == [ + None, + None, + None, + 4, + ] + with pytest.raises(click.ClickException): + _numpy_dtype_for(SimpleNamespace(name="BFLOAT16")) + + +def test_io_metadata_replaces_ordinal_names(tmp_path: Path) -> None: + model_path = tmp_path / "model.mlir" + metadata_path = tmp_path / "model_metadata.json" + metadata_path.write_text( + json.dumps( + { + "inputs": [{"name": "pixels", "index": 0}], + "outputs": [{"name": "scores", "index": 0}], + } + ) + ) + io_config = {"input_names": ["input_0"], "output_names": ["output_0"]} + + _apply_io_metadata(io_config, model_path) + + assert io_config["input_names"] == ["pixels"] + assert io_config["output_names"] == ["scores"] + + +def test_onnx_io_ranges_reach_tensor_comparison(tmp_path: Path) -> None: + from onnx import TensorProto, helper, save_model + + from winml.modelkit.eval.tensor_similarity_evaluator import TensorSimilarityEvaluator + from winml.modelkit.onnx import get_io_config + + model_path = tmp_path / "model.onnx" + graph = helper.make_graph( + [helper.make_node("Identity", ["indices"], ["output"])], + "input_ranges", + [helper.make_tensor_value_info("indices", TensorProto.INT64, ["batch", 16])], + [helper.make_tensor_value_info("output", TensorProto.INT64, ["batch", 16])], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + helper.set_model_props( + model, + {"winml.io.inputs": json.dumps([{"name": "indices", "value_range": [0, 1]}])}, + ) + save_model(model, model_path) + source_io = get_io_config(model_path) + io_config = { + "input_names": source_io["input_names"], + "input_shapes": [[2, 16]], + "input_types": ["int64"], + "output_names": source_io["output_names"], + } + + _apply_io_metadata(io_config, model_path) + + assert io_config["value_ranges"] == source_io["value_ranges"] + assert io_config["input_shapes"] == [[2, 16]] + assert io_config["input_types"] == ["int64"] + evaluator = object.__new__(TensorSimilarityEvaluator) + evaluator.model = SimpleNamespace(io_config=io_config) + evaluator.config = SimpleNamespace( + input_data=None, dataset=SimpleNamespace(samples=2, seed=42) + ) + + dataset = evaluator.prepare_data() + + lower, upper = source_io["value_ranges"]["indices"] + for sample in dataset: + indices = sample["indices"] + assert tuple(indices.shape) == tuple(io_config["input_shapes"][0]) + assert bool(((indices >= lower) & (indices < upper)).all()) + + +@pytest.mark.parametrize( + "failure_point", + ["_load_mlir", "_load_onnx_on_cgc", "_load_onnx_on_ort", + "create_pipeline_builder", "build", "_stage_diagnostics"], +) +def test_build_failure_releases_adapter(monkeypatch, mlir_target, failure_point): + ep_device, adapter = mlir_target + close = Mock(wraps=adapter.close) + monkeypatch.setattr(adapter, "close", close) + runtime = _Runtime(_Stage(), _Pipeline()) + wr = SimpleNamespace(Runtime=lambda: runtime, NotSupportedError=_NotSupportedError) + monkeypatch.setattr("winml.modelkit.session.runtime_session.import_runtime", lambda: wr) + session = WinMLRuntimeSession("model.mlir", ep_device=ep_device, backend="cgc") + if failure_point.startswith("_load_onnx"): + session._is_mlir = False + if failure_point == "_load_onnx_on_ort": + session._backend = "ort" + monkeypatch.setattr(session, "_resolve_target", lambda *_args: SimpleNamespace( + adapter=adapter, execution_target=adapter.pointer, + provider_name=None, device_class="gpu", + )) + failure = RuntimeError("adapter cleanup probe") + failing_call = Mock(side_effect=failure) + if failure_point.startswith("_load_"): + monkeypatch.setattr(session, failure_point, failing_call) + elif failure_point == "_stage_diagnostics": + monkeypatch.setattr( + "winml.modelkit.session.runtime_session._stage_diagnostics", failing_call, + ) + elif failure_point == "build": + monkeypatch.setattr(runtime.builder, failure_point, failing_call) + else: + monkeypatch.setattr(runtime, failure_point, failing_call) + with pytest.raises(RuntimeError, match="adapter cleanup probe") as raised: + session.compile() + assert raised.value is failure + close.assert_called_once_with() + assert session._adapter_handle is None + assert session._built is False + session.close() + close.assert_called_once_with() + + +@pytest.mark.parametrize("named_bindings", [False, True]) +def test_run_requests_all_outputs_before_each_execution(monkeypatch, mlir_target, named_bindings): + stage = _Stage() + pipeline = _Pipeline() + runtime = _Runtime(stage, pipeline) + wr = SimpleNamespace(Runtime=lambda: runtime, NotSupportedError=_NotSupportedError) + monkeypatch.setattr("winml.modelkit.session.runtime_session.import_runtime", lambda: wr) + names = ["sum", "product"] + named = SimpleNamespace( + bind_input=lambda _name, tensor: stage.bind_input(0, tensor), + output=lambda name: stage.output(names.index(name)), + ) + monkeypatch.setattr(stage, "ort_bindings", lambda: named, raising=False) + requests = [] + + def request_output(index): + requests.append(index) + stage.requested_outputs.add(index) + + def execute(): + assert stage.requested_outputs == set(range(len(names))) + pipeline.runs += 1 + + def output(index): + assert index in stage.requested_outputs + data = stage.bound[0].to_numpy() + return _Tensor(np.add(data, data) if index == 0 else np.multiply(data, data)) + + monkeypatch.setattr(stage, "request_output", request_output) + monkeypatch.setattr(stage, "output", output) + monkeypatch.setattr(pipeline, "run", execute) + session = WinMLRuntimeSession("model.mlir", ep_device=mlir_target[0], backend="cgc") + try: + session.compile() + session._has_named_bindings = named_bindings + session._io_config["output_names"] = names + random = np.random.default_rng(42) + with session.perf(warmup=0): + for _iteration in range(3): + data = random.normal(size=(2, 3, 8, 8)).astype(np.float32) + actual = session.run({"input_0": data}) + np.testing.assert_array_equal(actual["sum"], np.add(data, data)) + np.testing.assert_array_equal(actual["product"], np.multiply(data, data)) + assert requests == list(range(len(names))) * pipeline.runs + assert pipeline.runs == 3 + finally: + session.close() + + +def test_mlir_session_builds_runs_and_resets( + monkeypatch: pytest.MonkeyPatch, + mlir_target: tuple[SimpleNamespace, _AdapterHandle], +) -> None: + stage = _Stage() + pipeline = _Pipeline() + runtime = _Runtime(stage, pipeline) + ep_device, adapter = mlir_target + close = Mock(wraps=adapter.close) + monkeypatch.setattr(adapter, "close", close) + wr = SimpleNamespace( + Runtime=lambda: runtime, + NotSupportedError=_NotSupportedError, + ) + monkeypatch.setattr( + "winml.modelkit.session.runtime_session.import_runtime", + lambda: wr, + ) + + session = WinMLRuntimeSession("model.mlir", ep_device=ep_device, backend="cgc") + session.compile() + session.compile() + close.assert_not_called() + assert session._adapter_handle is adapter + assert runtime.builder.targets == [123] + assert runtime.builder.caller_owned_at_build is False + assert session.device == "gpu" + assert session.io_config["input_names"] == ["input_0"] + + outputs = session.run( + {"input_0": np.zeros((2, 3, 8, 8), dtype=np.float64)} + ) + assert outputs["output_0"].shape == (2, 5) + assert stage.bound[0].to_numpy().dtype == np.float32 + assert pipeline.runs == 1 + + session.reset() + assert session._pipeline is None + assert adapter.closed is True + close.assert_called_once_with() + session.close() + close.assert_called_once_with() + + +@pytest.mark.parametrize("dtype", [np.float32, np.float64]) +@pytest.mark.parametrize( + ("shape", "transpose"), + [((), False), ((1,), False), ((2, 3), False), ((2, 3), True)], + ids=["scalar", "vector", "matrix", "noncontiguous"], +) +def test_prepare_inputs_preserves_shape( + dtype: type, shape: tuple[int, ...], transpose: bool, +) -> None: + values = np.random.default_rng(0).standard_normal(shape).astype(dtype) + if transpose: + values = values.T + session = WinMLRuntimeSession("model.onnx", device="cpu", ep="cpu", backend="ort") + try: + session._io_config = {"input_names": ["input"], "input_types": [np.float32]} + + prepared = session._prepare_inputs({"input": values})["input"] + + assert prepared.shape == values.shape + assert prepared.dtype == np.float32 + assert prepared.flags.c_contiguous + np.testing.assert_array_equal(prepared, values.astype(np.float32)) + finally: + session.close() + + +def test_mlir_session_requires_ep_device() -> None: + with pytest.raises(ValueError, match="ep_device is required"): + WinMLRuntimeSession("model.mlir", backend="cgc") + + +@pytest.mark.parametrize("model_suffix", [".mlir", ".onnx"]) +@pytest.mark.parametrize("operation", ["compile", "run", "io_config"]) +def test_cgc_target_rejects_dxcore_incompatible_device( + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, + model_suffix: str, + operation: str, +) -> None: + projected = SimpleNamespace( + ep_name="DmlExecutionProvider", + device_type="GPU", + hardware_name="Projected GPU", + adapter_luid=141027, + ort_handle=object(), + ) + physical = SimpleNamespace( + ep_name="DmlExecutionProvider", + device_type="GPU", + hardware_name="Physical GPU", + adapter_luid=95733, + ) + physical_handle = object() + adapter = _AdapterHandle() + + monkeypatch.setattr( + "onnxruntime.get_ep_devices", + lambda: [physical_handle], + ) + monkeypatch.setattr( + "winml.modelkit.session.ep_device.WinMLDevice", + lambda handle: physical if handle is physical_handle else None, + ) + + attempted_luids = [] + native_error = click.ClickException("not visible to DXCore") + + def from_luid(cls: type[_DXCoreAdapter], luid: int) -> _AdapterHandle: + attempted_luids.append(luid) + if luid == 141027: + raise native_error + return adapter + + monkeypatch.setattr(_DXCoreAdapter, "from_luid", classmethod(from_luid)) + monkeypatch.setattr( + "winml.modelkit.session.runtime_session.import_runtime", + lambda: SimpleNamespace(Runtime=SimpleNamespace), + ) + + session = WinMLRuntimeSession( + f"model{model_suffix}", ep_device=SimpleNamespace(device=projected), backend="cgc" + ) + with pytest.raises(Exception) as exc_info: + if operation == "compile": + session.compile() + elif operation == "run": + session.run({"input": np.zeros((), dtype=np.float32)}) + else: + _ = session.io_config + + assert isinstance(exc_info.value, click.ClickException) + assert exc_info.value.__cause__ is native_error + assert "Cannot resolve the selected device LUID" in str(exc_info.value) + assert "winml sys" in str(exc_info.value) + assert "--device-luid " in str(exc_info.value) + assert attempted_luids == [projected.adapter_luid] + assert adapter.closed is False + assert "CGC is falling back" not in caplog.text + + +@pytest.mark.parametrize("missing_luid", [False, True]) +@pytest.mark.parametrize("model_suffix", [".mlir", ".onnx"]) +def test_cgc_target_uses_selected_adapter( + mlir_target: tuple[SimpleNamespace, _AdapterHandle], + missing_luid: bool, + model_suffix: str, +) -> None: + ep_device, adapter = mlir_target + device = ep_device.device + runtime = _Runtime(_Stage(), _Pipeline()) + session = WinMLRuntimeSession(f"model{model_suffix}", ep_device=ep_device, backend="cgc") + if missing_luid: + device.adapter_luid = None + with pytest.raises(Exception) as exc_info: + session._resolve_target(runtime, SimpleNamespace()) + assert isinstance(exc_info.value, click.ClickException) + assert exc_info.value.__cause__ is None + assert "Cannot resolve the selected device LUID" in str(exc_info.value) + assert "winml sys" in str(exc_info.value) + assert "--device-luid " in str(exc_info.value) + assert adapter.closed is False + else: + resolved = session._resolve_target(runtime, SimpleNamespace()) + try: + assert resolved.execution_target == adapter.pointer + assert resolved.adapter is adapter + assert adapter.closed is False + finally: + adapter.close() + + +def test_runtime_session_rejects_provider_options() -> None: + with pytest.raises(click.ClickException, match="--ep-options"): + WinMLRuntimeSession("model.onnx", provider_options={"key": "value"}, backend="ort") + + +def test_onnx_session_passes_resolved_ep_and_device_to_runtime( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + "winml.modelkit.onnx.get_io_config", lambda _path: {"value_ranges": {}} + ) + stage = _Stage() + pipeline = _Pipeline() + + class OrtModel(_Model): + def ort_schema(self) -> None: + return None + + class OrtRuntime(_Runtime): + def __init__(self) -> None: + super().__init__(stage, pipeline) + self.target_args: tuple[str, Any, Any] | None = None + + def load_model(self, path: str) -> _Model: + assert path.endswith(".onnx") + return OrtModel() + + def create_ort_execution_target( + self, + provider_name: str, + kind: Any, + hardware_target: Any, + ) -> object: + self.target_args = (provider_name, kind, hardware_target) + return object() + + runtime = OrtRuntime() + gpu_kind = object() + adapter = _AdapterHandle() + wr = SimpleNamespace( + Runtime=lambda: runtime, + NotSupportedError=_NotSupportedError, + ExecutionTargetKind=SimpleNamespace(GPU=gpu_kind), + ) + ep_device = SimpleNamespace( + device=SimpleNamespace( + ep_name="NvTensorRTRTXExecutionProvider", + device_type="GPU", + hardware_name="Test GPU", + adapter_luid=456, + ), + ep_short_name="nvtensorrtrtx", + source_tag="winml-catalog", + ) + monkeypatch.setattr( + "winml.modelkit.session.runtime_session.import_runtime", + lambda: wr, + ) + monkeypatch.setattr( + "winml.modelkit.session.runtime_session.resolve_provider_kind", + lambda *_args: (_ for _ in ()).throw( + AssertionError("concrete ep_device must bypass request-based resolution") + ), + ) + monkeypatch.setattr( + _DXCoreAdapter, + "from_luid", + classmethod(lambda cls, luid: adapter if luid == 456 else None), + ) + + session = WinMLRuntimeSession("model.onnx", ep_device=ep_device, backend="ort") + session.compile() + + assert runtime.target_args == ( + "NvTensorRTRTXExecutionProvider", + gpu_kind, + 123, + ) + assert runtime.builder.caller_owned_at_build is False + assert runtime.builder.targets[0] is not None + assert session.device == "gpu" + assert session.requested_provider == "NvTensorRTRTXExecutionProvider" + session.reset() + assert adapter.closed is True + + +def test_onnx_session_without_ep_device_uses_request_resolution( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + "winml.modelkit.onnx.get_io_config", lambda _path: {"value_ranges": {}} + ) + stage = _Stage() + pipeline = _Pipeline() + + class OrtModel(_Model): + def ort_schema(self) -> None: + return None + + class OrtRuntime(_Runtime): + def __init__(self) -> None: + super().__init__(stage, pipeline) + self.target_args: tuple[str, Any] | None = None + + def load_model(self, path: str) -> _Model: + assert path.endswith(".onnx") + return OrtModel() + + def create_ort_execution_target(self, provider_name: str, kind: Any) -> object: + self.target_args = (provider_name, kind) + return object() + + runtime = OrtRuntime() + gpu_kind = object() + wr = SimpleNamespace( + Runtime=lambda: runtime, + NotSupportedError=_NotSupportedError, + ExecutionTargetKind=SimpleNamespace(GPU=gpu_kind), + ) + calls: list[tuple[Any, ...]] = [] + + def resolve(*args: Any) -> tuple[str, str]: + calls.append(args) + return "DmlExecutionProvider", "gpu" + + monkeypatch.setattr( + "winml.modelkit.session.runtime_session.import_runtime", + lambda: wr, + ) + monkeypatch.setattr( + "winml.modelkit.session.runtime_session.resolve_provider_kind", + resolve, + ) + + session = WinMLRuntimeSession( + "model.onnx", + device="gpu", + ep="dml", + backend="ort", + ) + session.compile() + + assert calls == [("gpu", "dml", None)] + assert runtime.target_args == ("DmlExecutionProvider", gpu_kind) + + +def test_concurrent_compile_builds_once( + monkeypatch: pytest.MonkeyPatch, + mlir_target: tuple[SimpleNamespace, _AdapterHandle], +) -> None: + stage = _Stage() + pipeline = _Pipeline() + build_entered = threading.Event() + release_build = threading.Event() + + class BlockingBuilder(_Builder): + builds = 0 + + def build(self) -> _Pipeline: + self.builds += 1 + build_entered.set() + assert release_build.wait(timeout=5) + return self.pipeline + + runtime = _Runtime(stage, pipeline) + runtime.builder = BlockingBuilder(stage, pipeline) + wr = SimpleNamespace(Runtime=lambda: runtime, NotSupportedError=_NotSupportedError) + monkeypatch.setattr( + "winml.modelkit.session.runtime_session.import_runtime", + lambda: wr, + ) + ep_device, _ = mlir_target + session = WinMLRuntimeSession("model.mlir", ep_device=ep_device, backend="cgc") + first = threading.Thread(target=session.compile) + second = threading.Thread(target=session.compile) + + first.start() + assert build_entered.wait(timeout=5) + second.start() + time.sleep(0.05) + assert runtime.builder.builds == 1 + release_build.set() + first.join(timeout=5) + second.join(timeout=5) + + assert not first.is_alive() + assert not second.is_alive() + assert runtime.builder.builds == 1 + + +def test_concurrent_runs_serialize_bind_run_read( + monkeypatch: pytest.MonkeyPatch, + mlir_target: tuple[SimpleNamespace, _AdapterHandle], +) -> None: + run_entered = threading.Event() + release_run = threading.Event() + + class TrackingStage(_Stage): + def __init__(self) -> None: + super().__init__() + self.bind_count = 0 + + def bind_input(self, index: int, tensor: _Tensor) -> None: + self.bind_count += 1 + super().bind_input(index, tensor) + + def output(self, index: int) -> _Tensor: + value = float(self.bound[index].to_numpy().flat[0]) + return _Tensor(np.full((1, 5), value, dtype=np.float32)) + + class BlockingPipeline(_Pipeline): + def run(self) -> None: + self.runs += 1 + if self.runs == 1: + run_entered.set() + assert release_run.wait(timeout=5) + + stage = TrackingStage() + pipeline = BlockingPipeline() + runtime = _Runtime(stage, pipeline) + wr = SimpleNamespace(Runtime=lambda: runtime, NotSupportedError=_NotSupportedError) + monkeypatch.setattr( + "winml.modelkit.session.runtime_session.import_runtime", + lambda: wr, + ) + ep_device, _ = mlir_target + session = WinMLRuntimeSession("model.mlir", ep_device=ep_device, backend="cgc") + outputs: dict[str, float] = {} + + def run(name: str, value: float) -> None: + result = session.run( + {"input_0": np.full((1, 3, 8, 8), value, dtype=np.float32)} + ) + outputs[name] = float(result["output_0"][0, 0]) + + first = threading.Thread(target=run, args=("first", 1.0)) + second = threading.Thread(target=run, args=("second", 2.0)) + first.start() + assert run_entered.wait(timeout=5) + second.start() + time.sleep(0.05) + assert stage.bind_count == 1 + release_run.set() + first.join(timeout=5) + second.join(timeout=5) + + assert outputs == {"first": 1.0, "second": 2.0} + + +def test_close_waits_for_active_run( + monkeypatch: pytest.MonkeyPatch, + mlir_target: tuple[SimpleNamespace, _AdapterHandle], +) -> None: + run_entered = threading.Event() + release_run = threading.Event() + + class BlockingPipeline(_Pipeline): + def __init__(self) -> None: + super().__init__() + self.closed = False + + def run(self) -> None: + run_entered.set() + assert release_run.wait(timeout=5) + + def close(self) -> None: + self.closed = True + + stage = _Stage() + pipeline = BlockingPipeline() + runtime = _Runtime(stage, pipeline) + wr = SimpleNamespace(Runtime=lambda: runtime, NotSupportedError=_NotSupportedError) + monkeypatch.setattr( + "winml.modelkit.session.runtime_session.import_runtime", + lambda: wr, + ) + ep_device, _ = mlir_target + session = WinMLRuntimeSession("model.mlir", ep_device=ep_device, backend="cgc") + inference = threading.Thread( + target=session.run, + args=({"input_0": np.zeros((1, 3, 8, 8), dtype=np.float32)},), + ) + teardown = threading.Thread(target=session.close) + + inference.start() + assert run_entered.wait(timeout=5) + teardown.start() + time.sleep(0.05) + assert pipeline.closed is False + release_run.set() + inference.join(timeout=5) + teardown.join(timeout=5) + + assert not inference.is_alive() + assert not teardown.is_alive() + assert pipeline.closed is True + + +def test_perf_rejects_monitor_without_loading_runtime() -> None: + session = WinMLRuntimeSession( + "model.mlir", ep_device=_mlir_ep_device(), backend="cgc" + ) + with pytest.raises(click.ClickException, match="monitor"), session.perf(monitor=object()): + pass diff --git a/tests/unit/session/test_winml_device.py b/tests/unit/session/test_winml_device.py index a1d60095f..a96434ffd 100644 --- a/tests/unit/session/test_winml_device.py +++ b/tests/unit/session/test_winml_device.py @@ -91,6 +91,18 @@ def test_hardware_name_unknown_fallback(self) -> None: handle = make_fake_ort_ep_device(ep_name="OpenVINOExecutionProvider", device_type="NPU") assert WinMLDevice(handle).hardware_name == "" + def test_adapter_luid_reads_device_metadata(self) -> None: + handle = make_fake_ort_ep_device( + ep_name="OpenVINOExecutionProvider", + device_type="GPU", + device_metadata={"LUID": "57733"}, + ) + assert WinMLDevice(handle).adapter_luid == 57733 + + def test_adapter_luid_missing_returns_none(self) -> None: + handle = make_fake_ort_ep_device(ep_name="OpenVINOExecutionProvider", device_type="CPU") + assert WinMLDevice(handle).adapter_luid is None + def test_vendor_passes_through(self) -> None: handle = make_fake_ort_ep_device( ep_name="OpenVINOExecutionProvider", diff --git a/tests/unit/utils/test_cli.py b/tests/unit/utils/test_cli.py index 51b9c5732..38c50c75f 100644 --- a/tests/unit/utils/test_cli.py +++ b/tests/unit/utils/test_cli.py @@ -8,6 +8,7 @@ import json import os +from dataclasses import dataclass from typing import TYPE_CHECKING import click @@ -29,6 +30,7 @@ optimize_option, overwrite_option, parse_ep_options, + parse_options, precision_option, quant_option, ) @@ -82,6 +84,71 @@ def test_empty_key_raises(self) -> None: parse_ep_options(("=value",)) +class TestParseOptions: + """Tests for dataclass-driven primitive KEY=VALUE options.""" + + @dataclass(frozen=True) + class Options: + enabled: bool = False + count: int = 1 + ratio: float = 1.0 + label: str = "default" + + def test_parses_primitive_values(self) -> None: + assert parse_options( + ( + "enabled=true", + "count=4", + "ratio=1.5", + "label=fast", + ), + self.Options, + ) == self.Options(enabled=True, count=4, ratio=1.5, label="fast") + + def test_false_is_case_insensitive(self) -> None: + assert parse_options( + ("enabled=FALSE",), + self.Options, + ) == self.Options() + + def test_duplicate_key_uses_last_value(self) -> None: + assert parse_options( + ("count=1", "count=2"), + self.Options, + ) == self.Options(count=2) + + def test_hyphenated_key_maps_to_underscored_field(self) -> None: + @dataclass(frozen=True) + class HyphenatedOptions: + external_weights: bool = False + + assert parse_options( + ("external-weights=true",), + HyphenatedOptions, + ) == HyphenatedOptions(external_weights=True) + + def test_empty_values_use_dataclass_defaults(self) -> None: + assert parse_options((), self.Options) == self.Options() + + @pytest.mark.parametrize( + "value", + [ + "missing-separator", + "=true", + "unknown=true", + "enabled=yes", + "count=one", + "ratio=fast", + ], + ) + def test_invalid_value_raises(self, value: str) -> None: + with pytest.raises(click.BadParameter): + parse_options( + (value,), + self.Options, + ) + + class TestNoColorOption: """Tests for the shared no_color_option() decorator.""" diff --git a/tests/unit/utils/test_runtime_constants.py b/tests/unit/utils/test_runtime_constants.py index 2962220e2..dea1dd514 100644 --- a/tests/unit/utils/test_runtime_constants.py +++ b/tests/unit/utils/test_runtime_constants.py @@ -8,9 +8,46 @@ from typing import get_args -from winml.modelkit.utils.constants import RUNTIME_NAMES, RuntimeName +import pytest + +from winml.modelkit.utils.constants import ( + RUNTIME_BACKENDS, + RUNTIME_NAMES, + RuntimeBackend, + RuntimeName, + resolve_runtime_api_backend, +) def test_runtime_names_match_runtime_name_literal() -> None: assert get_args(RuntimeName) == RUNTIME_NAMES - assert RUNTIME_NAMES == ("auto", "winml-ort", "ort-genai") + assert RUNTIME_NAMES == ("auto", "winml-ort", "ort-genai", "winml-runtime") + + +def test_runtime_backends_match_runtime_backend_literal() -> None: + assert get_args(RuntimeBackend) == RUNTIME_BACKENDS + assert RUNTIME_BACKENDS == ("ort", "cgc") + + +@pytest.mark.parametrize( + ("model_path", "backend", "expected"), + [ + ("model.onnx", None, "cgc"), + ("model.onnx", "ort", "ort"), + ("model.onnx", "cgc", "cgc"), + ("model.mlir", None, "cgc"), + ("model.mlir", "cgc", "cgc"), + ], +) +def test_resolve_runtime_api_backend(model_path, backend, expected) -> None: + assert resolve_runtime_api_backend("winml-runtime", model_path, backend) == expected + + +def test_resolve_runtime_api_backend_rejects_ort_for_mlir() -> None: + with pytest.raises(ValueError, match="MLIR inputs require"): + resolve_runtime_api_backend("winml-runtime", "model.mlir", "ort") + + +def test_resolve_runtime_api_backend_rejects_other_runtimes() -> None: + with pytest.raises(ValueError, match="only supported"): + resolve_runtime_api_backend("winml-ort", "model.onnx", "cgc")