From 0091a66bb66096981e71ad3b9f68508a7f40777c Mon Sep 17 00:00:00 2001 From: fangyangci <133664123+fangyangci@users.noreply.github.com> Date: Sun, 20 Sep 2026 15:03:12 +0800 Subject: [PATCH 1/6] chore(release): prepare v0.4.0 (#1431) ## Summary - Add the v0.4.0 CHANGELOG entry from the reviewed monthly release scope. - Bump the package version from 0.3.1 to 0.4.0 without changing runtime dependencies. Targets release/v0.4.0 for the planned 2026-09-24 release. CHANGELOG and version changes are separate commits; the ignored local uv.lock is not included. ## Validation - Configured pre-commit checks passed for both files. - uv lock --check passed using the approved index, and uv version reports 0.4.0. - Release-note PR references and git diff --check passed. Final release E2E, model gates, official-wheel smoke, and publication approvals remain separate. --- CHANGELOG.md | 68 ++++++++++++++++++++++++++++++++++++++++++++++++++ pyproject.toml | 2 +- 2 files changed, 69 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 3836542ae..985af62c4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,74 @@ All notable changes to this project are documented in this file. The format is loosely based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/). +## WinML CLI v0.4.0 + +This cycle adds CGIR and Windows ML Runtime workflows, native PyTorch evaluation, physical-adapter selection, and broader operator profiling. It improves pattern-aware analysis, targeted optimization, model-loading memory measurements, and release evaluation. The release also carries forward the fixes already published in v0.3.1. See the behavior changes below. + +### ⚠️ Behavior changes + +- **Runtime selection** - runtime names are now `winml-ort` and `ort-genai` instead of `winml` and `winml-genai`; `winml perf` defaults to automatic runtime selection while retaining explicit choices (#1366, #1308). +- `winml perf` - canonical cache controls are `--use-cache` / `--no-use-cache`; the older `--ignore-cache` / `--no-ignore-cache` forms remain hidden deprecated aliases, and conflicting toggles are rejected (#1271). +- `winml analyze` - removes the obsolete `--htp-metadata` option in favor of pattern-rule-based runtime support checks (#1218). +- `winml serve` - replaces wildcard CORS with same-origin request protection, validates allowed CLI commands, and rejects HTTP attempts to enable remote code, including through nested model/configuration loading (#1321). +- **Memory reports** - classic `memory_measurement.schema_version` is now `3`, with revised baseline/delta boundaries; the outer report remains schema version `2`, and load-only measurements use `load_memory.version = 1` (#1429). +- **Operator tracing** - measured tracing runs default to 10 iterations; explicit iteration overrides remain available (#1406). +- **Evaluation recipe labels** - recipes previously labelled FP16 without an actual quantization/conversion step are labelled FP32; historical result paths and reference labels are preserved rather than rewritten (#1411). + +### ✨ Improvements + +- **CGIR workflows** - `export`, `config`, `build`, `compile`, `perf`, and `eval` integrate standalone MLIR, Windows ML Runtime with CGC compilation, and ONNX through the WinMLCG EP, using the required preview Runtime/EP packages (#1426). +- `winml eval --runtime pytorch` - evaluates Hugging Face models natively on CPU or CUDA and accepts existing PyTorch models through the evaluation API, preserving checkpoint dtype and reporting the selected runtime (#1282). +- `winml eval` - adds independent `--reference-device` and `--reference-ep` controls for ONNX comparison and integrates the shared device-selection path (#1324, #1412). +- `winml perf --device-luid` - selects the physical adapter reported by `winml sys`, keeps inference and monitoring on that adapter, and distinguishes identical GPUs; the option does not apply to `ort-genai` (#1398). +- **Operator profiling** - adds basic OpenVINO CPU/NPU and TensorRT RTX GPU tracing, excludes warmup samples, preserves raw trace artifacts, and improves QNN schematics and detail-fallback diagnostics (#1377, #1406, #1288, #1289). +- **Model-loading memory** - adds load-only measurements from runtime/device readiness through model compilation, separates sampled peaks from OS lifetime diagnostics, and prepares missing process counters on the explicitly selected GPU while preserving unavailable values and signed deltas (#1429). +- `winml sys` - uses DXCore for GPU/NPU identity and adapter LUIDs, retains WMI/PnP enrichment and fallback, and recognizes NVIDIA ACPI PnP identifiers (#1351, #1290). +- `winml sys` - expands memory metadata, dedicated/shared GPU memory, NPU memory, and Windows build details in system reports (#1312, #1387, #1421, #1425). +- `winml analyze` - evaluates matched subgraphs with pattern-level runtime rules before operator-level fallback, reports pattern coverage, and includes optimization findings in JSON output (#1218, #1285). +- `winml optimize` - adds EP/device targeting, displays custom operator domains, and accelerates capability checks (#1256, #1306, #1298). +- **Graph optimization** - adds routed affine and positive-Exp scale folding while preserving fast QNN grouped-convolution regions (#1301, #1317). +- **Qwen3 GenAI bundles** - shares decoder context/iterator weights through a common EP context and adds VitisAI support for the transformer stages (#1305). +- **Model caches** - unifies cache/rebuild controls across build, eval, and perf; evaluation enables model-cache reuse by default and reports controls that do not apply to prebuilt inputs (#1269, #1270, #1271). +- `winml export --batch-size` - supports validated static batch sizes and carries them through input specifications and export metadata (#1315). +- **Model adapters** - adds the Unlimited-OCR vision tower for feature extraction and audeering Wav2Vec2 dimensional-emotion speech regression (#1018, #1084). +- **Vision recipes** - adds CPU configurations for RT-DETR, ViTPose, and OWLv2 zero-shot detection, plus QNN NPU LayoutLM document-QA configurations (#1190, #1189, #1196, #1296, #1369). +- **Language recipes** - adds mMARCO MiniLM, NLI MiniLM, multilingual E5, and Spanish BERT configurations (#1191, #1192, #1210, #1185). +- **Audio recipes** - adds Wav2Vec2 deepfake classification, MMS-1B-all CPU configurations, and a QNN NPU W8A16 configuration for dimensional-emotion regression (#1194, #1177, #1318). + +### 🐛 Fixes + +- **Hugging Face export** - repairs Marian, TrOCR/Manga-OCR, LayoutLM QA, BLIP decoder, and SAM wrapper paths, and corrects eager-attention selection during model loading (#1323, #1330, #1392, #1403, #1375). +- **FP16 conversion** - handles initializer-backed graph outputs, prevents quantization naming collisions, and validates captured local-function tensors against stricter ONNX Runtime graph-attribute requirements (#1280, #1417, #1377). +- **Calibration and quantization** - clamps DistilBERT mask constants before calibration and fixes input-dtype, synthetic-calibration, and large-model QDQ serialization regressions found during release evaluation (#1420, #1411). +- **FP16 recipes** - restores missing quantization configuration for RoBERTa-large SQuAD2 and DeBERTa-v3-base NLI recipes (#1245, #1244). +- **ViTPose export** - derives dummy inputs from the model configuration (#1299). +- **Compiled-model caches** - keys EPContext reuse by compile identity and handles multiple QNN EP context partitions in perf (#1295, #1361). +- **GenAI perf** - honors the selected device and EP during automatic builds and aligns performance metric schemas (#1404, #1307). +- **Performance monitoring** - corrects multi-GPU monitoring, preserves memory baselines, and distinguishes unavailable measurements from real samples (#1313, #1414). +- **Runtime initialization** - includes CPU in vendor compatibility, lazily loads session backends and monitors, and removes the manual ONNX Runtime DLL preload path (#1303, #1382, #1380). +- `winml analyze` - repairs schema fallback for runtime-specific operators (#1304). +- **Image-to-text evaluation** - repairs evaluator behavior (#1410). +- `winml sys --format json` - keeps EP installation notices and download progress out of JSON output while retaining them in human-readable modes (#1314). + +### 🔧 Internals & CI + +- **Release evaluation** - adds an opt-in manifest-driven release sweep with explicit model/task, precision, and machine/EP/device targeting, while preserving non-release selection modes; cross-EP evaluation and structured perf-result handling are improved (#1411, #1281, #1316). +- **Evaluation reliability** - adds targeted timeout exclusions and failure categorization, separates download time from execution budgets, prioritizes model execution, and prevents stage logs from being mistaken for final artifact paths (#1358, #1357, #1400, #1427). +- **LLM evaluation** - adds a schema-normalized evaluation runner (#1277). +- **E2E stability** - isolates native EP CLI invocations in subprocesses, improves shared-RDP GPU coverage, and updates memory assertions to the current RAM label (#1416, #1418, #1430). +- **GitHub releases** - uses the GitHub App service connection and preserves UTF-8 release notes, including BOM-aware output for release tasks (#1292, #1293). +- **CI supply chain** - pins GitHub Actions, updates action dependencies, and introduces a seven-day Dependabot cooldown (#1408, #1409). +- **Dependency maintenance** - updates aiohttp, cryptography, Jupyter dependencies, and the Windows ML ONNX Runtime compatibility range; removes unused Jupyter server packages from the development dependency set (#1278, #1279, #1216, #1386, #1423). +- **Runtime dispatch** - centralizes perf runtime names in shared constants (#1345). +- **Contributor workflows** - adds skills for model-support contributions and correctness-gated auto-optimization (#1415, #1428). +- **Release-line synchronization** - carries the v0.3.0 and v0.3.1 release changes back to main, including the dataset-ID and native-warning-spooling fixes already shipped through earlier release cherry-picks (#1291, #1397, #1262, #1266). + +### 📦 Assets + +- `winml_cli-0.4.0-py3-none-any.whl` +- `rules-v0.4.0.zip` + ## WinML CLI v0.3.1 This hotfix restores compatibility with current Transformers and Windows runtime dependencies, fixes VitisAI cache permissions, and stabilizes Hugging Face model export and evaluation. diff --git a/pyproject.toml b/pyproject.toml index a8763a046..89d7de81b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ requires = [ "setuptools>=61", "wheel" ] [project] name = "winml-cli" -version = "0.3.1" +version = "0.4.0" description = "Accelerate Model Deployment on WinML" readme = "README.md" keywords = [ "onnx", "winml" ] From 855fe5c4dfaf2f501d4aacdfb0460939a813c26e Mon Sep 17 00:00:00 2001 From: Zhenchao Ni Date: Mon, 21 Sep 2026 15:27:16 +0800 Subject: [PATCH 2/6] Remove pad and identity rewrite (#1432) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary Remove two CGC compatibility workarounds following upstream IX fixes already included in the current compiler version (`35c01c26`). | Removed workaround | Upstream issue | Fix | |---|---|---| | `eliminate-identity` | microsoft/ix#1198 — ONNX Identity fails to lower to Foundry and DXCGC | microsoft/ix#1201 adds native Identity lowering. | | Pad-specific constant folding (`fold_constant_pad_pads`) and the `fold-constant-pad-pads` compatibility alias | microsoft/ix#1199 — Support constant folding of ONNX Pad parameters before lowering | microsoft/ix#1206 improves Cast/Pad chain support, static shape inference, and lowering type compatibility. | ## Changes - Remove the Identity rewrite implementation, capability registration, and public export. - Remove the Pad-specific folding prepass and its dedicated dead-producer cleanup. - Remove the `fold-constant-pad-pads` alias and public function export. - Update documentation and tests to reflect the retired rules. - Preserve `cgc-constant-folding` and its shared integer/boolean expression evaluator. General folding remains necessary for static shape chains and continues to fold Pad parameters in graphs containing `Shape`. - Leave compiler options, including the default-enabled topological sorting, unchanged. ## Validation - 203 affected unit tests passed. - All five models that previously used the Pad-specific prepass retained identical Pad parameter values. - GPU outputs before and after Pad-prepass removal were elementwise identical across three input samples per model. - The three SOD models required topological sorting to be disabled for the GPU comparison because the current wheel does not yet include the fix for microsoft/ix#1207. Their default-path conversion failure remains unchanged. - End-to-end SOD validation through JSON config, build, perf, and eval is still in progress. --- docs/commands/optimize.md | 3 +- .../optim/pipes/cgir_rewrite_rules.py | 26 +- src/winml/modelkit/pattern/cgc/__init__.py | 5 +- .../pattern/cgc/cgc_constant_folding.py | 153 +--------- .../modelkit/pattern/cgc/identity_rewrites.py | 206 ------------- .../optim/pipes/test_pipe_cgir_rewrite.py | 278 ++---------------- 6 files changed, 42 insertions(+), 629 deletions(-) delete mode 100644 src/winml/modelkit/pattern/cgc/identity_rewrites.py diff --git a/docs/commands/optimize.md b/docs/commands/optimize.md index 5c7134451..950fa8ec3 100644 --- a/docs/commands/optimize.md +++ b/docs/commands/optimize.md @@ -135,10 +135,9 @@ the shared pattern package does not re-export these backend-specific patterns. |------------|--------------------------| | `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. | +| `cgc-constant-folding` | Fill FoundryToolbox constant-folding gaps without an ORT Session. In graphs containing `Shape`, 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. Graphs without `Shape`, 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. | | `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. | diff --git a/src/winml/modelkit/optim/pipes/cgir_rewrite_rules.py b/src/winml/modelkit/optim/pipes/cgir_rewrite_rules.py index 6e5e951ef..475478e14 100644 --- a/src/winml/modelkit/optim/pipes/cgir_rewrite_rules.py +++ b/src/winml/modelkit/optim/pipes/cgir_rewrite_rules.py @@ -29,7 +29,6 @@ ResizeWithTfHalfPixelForNNPattern, cgc_constant_folding, deduplicate_opset_imports, - eliminate_identity, normalize_int32_dq, ) from ..registry import BoolCapability, CapabilityCategory @@ -73,26 +72,10 @@ class CGIRModelRewriteRule: 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", + description="Fill FoundryToolbox folding gaps for static shape subgraphs", category=CapabilityCategory.REWRITE, default=False, ) @@ -169,7 +152,6 @@ class CGIRModelRewriteRule: CGIRModelRewriteRule( capability=CGC_CONSTANT_FOLDING, transform=cgc_constant_folding, - aliases=(FOLD_CONSTANT_PAD_PADS,), ), CGIRModelRewriteRule( capability=DEDUPLICATE_OPSET_IMPORTS, @@ -214,10 +196,6 @@ class CGIRModelRewriteRule: target=MatMulDFTPattern, minimum_opset=17, ), - CGIRModelRewriteRule( - capability=ELIMINATE_IDENTITY, - transform=eliminate_identity, - ), CGIRRewriteRule( capability=GRIDSAMPLE_TO_GATHER, source=LinearGridSamplePattern, @@ -240,8 +218,6 @@ class CGIRModelRewriteRule: "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", diff --git a/src/winml/modelkit/pattern/cgc/__init__.py b/src/winml/modelkit/pattern/cgc/__init__.py index 82eda35ad..6aba322a5 100644 --- a/src/winml/modelkit/pattern/cgc/__init__.py +++ b/src/winml/modelkit/pattern/cgc/__init__.py @@ -4,12 +4,11 @@ # -------------------------------------------------------------------------- """Opt-in CGC compatibility patterns and model metadata rewrites.""" -from .cgc_constant_folding import cgc_constant_folding, fold_constant_pad_pads +from .cgc_constant_folding import cgc_constant_folding 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 ( @@ -39,7 +38,5 @@ "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 index f4cf2c49a..021cea034 100644 --- a/src/winml/modelkit/pattern/cgc/cgc_constant_folding.py +++ b/src/winml/modelkit/pattern/cgc/cgc_constant_folding.py @@ -3,7 +3,7 @@ # Licensed under the MIT License. # -------------------------------------------------------------------------- -"""Bounded constant folding for CGC Pad parameters and static shape subgraphs. +"""Bounded constant folding for CGC 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 @@ -14,13 +14,10 @@ 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, @@ -62,7 +59,6 @@ def __init__(self, model: ModelProto, *, static_shapes: bool = False) -> None: 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: @@ -153,7 +149,7 @@ def evaluate(self, name: str, visited: set[str], active: set[str]) -> np.ndarray raise ValueError("Only integer Cast targets are supported") fragment = helper.make_model( helper.make_graph( - [node], "constant_pad_parameter", [], + [node], "constant_parameter", [], [ValueInfoProto(name=name)], [numpy_helper.from_array(value, key) for key, value in inputs.items()], ), @@ -161,7 +157,6 @@ def evaluate(self, name: str, visited: set[str], active: set[str]) -> np.ndarray 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: @@ -173,150 +168,22 @@ def evaluate(self, name: str, visited: set[str], active: set[str]) -> np.ndarray 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. + """Fill FoundryToolbox constant-folding gaps for 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 + Only the main graph is changed. Graphs containing Shape fold bounded integer/boolean 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"}} + for node in model.graph.node): + return model + versions = {item.version for item in model.opset_import if item.domain in {"", "ai.onnx"}} if len(versions) != 1 or next(iter(versions)) < 11: - return prepared + return model rewritten = ModelProto() - rewritten.CopyFrom(prepared) + rewritten.CopyFrom(model) changed = False for _iteration in range(32): for value_info in rewritten.graph.value_info: @@ -339,7 +206,7 @@ def cgc_constant_folding(model: ModelProto) -> ModelProto: )) folded += 1 if not folded: - return rewritten if changed else prepared + return rewritten if changed else model 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") diff --git a/src/winml/modelkit/pattern/cgc/identity_rewrites.py b/src/winml/modelkit/pattern/cgc/identity_rewrites.py deleted file mode 100644 index 665f949ab..000000000 --- a/src/winml/modelkit/pattern/cgc/identity_rewrites.py +++ /dev/null @@ -1,206 +0,0 @@ -# ------------------------------------------------------------------------- -# 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/tests/unit/optim/pipes/test_pipe_cgir_rewrite.py b/tests/unit/optim/pipes/test_pipe_cgir_rewrite.py index 6d9a578c2..60c0212c6 100644 --- a/tests/unit/optim/pipes/test_pipe_cgir_rewrite.py +++ b/tests/unit/optim/pipes/test_pipe_cgir_rewrite.py @@ -129,19 +129,26 @@ def test_normalize_int32_dq_preserves_unsupported(guard): @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 +@pytest.mark.parametrize("has_shape", [False, True]) +def test_cgc_constant_folding_pad_parameters(opset, shared, has_shape): + from winml.modelkit.pattern.cgc import cgc_constant_folding seed = np.random.default_rng(42) widths = seed.integers(0, 3, size=4, dtype=np.int32) + nodes = [ + helper.make_node("Cast", ["widths"], ["pads"], to=TensorProto.INT64), + helper.make_node("Pad", ["source", "pads"], ["result"]), + ] outputs = [helper.make_tensor_value_info("result", TensorProto.FLOAT, [None, None])] + if has_shape: + nodes.append(helper.make_node("Shape", ["result"], ["shape"])) + outputs.append(helper.make_tensor_value_info("shape", TensorProto.INT64, [2])) 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])], + nodes, "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, @@ -150,14 +157,22 @@ def test_fold_constant_pad_pads(opset, shared): 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), + model, CGIRRewritePipe.build_config(cgc_constant_folding=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 + assert [(value.name, value.type.tensor_type.elem_type) for value in result.graph.output] == [ + (value.name, value.type.tensor_type.elem_type) for value in model.graph.output + ] + if has_shape: + for output, reference in zip(result.graph.output, expected, strict=True): + assert tuple(dim.dim_value for dim in output.type.tensor_type.shape.dim) == ( + reference.shape + ) + assert (result is model) == (not has_shape) + assert sum(node.op_type == "Cast" for node in result.graph.node) == int(not has_shape) + assert cgc_constant_folding(result) is result for actual, reference in zip( ReferenceEvaluator(result).run(None, feeds), expected, strict=True, ): @@ -234,141 +249,12 @@ def test_gridsample_to_gather_preserves_unsupported_modes(mode, padding): 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(): +def test_cgc_constant_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, - ) + assert len(CGIRRewritePipe.build_config(cgc_constant_folding=True).rules) == 1 @pytest.mark.parametrize("dynamic", [False, True]) @@ -480,7 +366,7 @@ def match(self): 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(CGIRRewritePipe.build_config(deduplicate_opset_imports=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)} @@ -506,6 +392,8 @@ def match(self): "log-to-reduce-log-sum", "materialize-initializer-parameters", "fold-scalar-initializer-casts", + "eliminate-identity", + "fold-constant-pad-pads", ]) def test_retired_cgir_rules_are_not_registered(capability): from click.testing import CliRunner @@ -522,114 +410,6 @@ def test_retired_cgir_rules_are_not_registered(capability): 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, From ff977dd98c68d6f1f520a69ea33e99cb4279cb86 Mon Sep 17 00:00:00 2001 From: Zhenchao Ni Date: Tue, 22 Sep 2026 13:53:55 +0800 Subject: [PATCH 3/6] Turn off toposort by default (#1434) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Disable topological sorting by default for CGIR conversion to work around microsoft/ix#1207, whose fix is not included in the compiler shipped with the current Windows ML wheel. This avoids conversion failures in the three P1 models and aligns with IX OnnxConverter’s default. Sorting remains available through `topo-sort-nodes=true`. --- src/winml/modelkit/export/cgc/exporter.py | 2 +- tests/unit/commands/test_export.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/winml/modelkit/export/cgc/exporter.py b/src/winml/modelkit/export/cgc/exporter.py index 6518517f9..55d8e4ceb 100644 --- a/src/winml/modelkit/export/cgc/exporter.py +++ b/src/winml/modelkit/export/cgc/exporter.py @@ -31,7 +31,7 @@ class CGCOptions: """Configuration shared by CGC export steps.""" external_weights: bool = False - topo_sort_nodes: bool = True + topo_sort_nodes: bool = False update_opset: bool = True freeze_dims: str = "" diff --git a/tests/unit/commands/test_export.py b/tests/unit/commands/test_export.py index a9eda7b73..6a614e283 100644 --- a/tests/unit/commands/test_export.py +++ b/tests/unit/commands/test_export.py @@ -151,7 +151,7 @@ def test_target_config_and_cli_option_precedence( 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}})) + config_path.write_text(json.dumps({"options": {"topo_sort_nodes": True}})) args += ["-c", str(build_path), "--export-config", str(config_path)] else: args += [config_flag, str(config_path)] @@ -178,7 +178,7 @@ def test_target_config_and_cli_option_precedence( 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", + topo_sort_nodes=not override and config_flag == "both", ) if not onnx_input: config = backend.call_args.kwargs["export_config"] From 0ce1251d58024d7c449a1d1a10e41f7160819255 Mon Sep 17 00:00:00 2001 From: fangyangci <133664123+fangyangci@users.noreply.github.com> Date: Tue, 22 Sep 2026 15:23:13 +0800 Subject: [PATCH 4/6] fix(packaging): bundle provider config for v0.4.0 (#1436) ## Summary - Bundle `analyze/utils/avalizble_ep_device_ops/avaliable_providers.json` in the wheel and source distribution so installed packages can resolve valid EP/device pairs. - Add a packaging regression test that builds the distributions and verifies the JSON is present with the expected contents in both archives. Targets `release/v0.4.0`. No runtime code, dependency, version, or pipeline changes. ## Validation - Focused wheel/sdist regression test: 1 passed on the latest release branch baseline. - Ruff check, Ruff format check, and `git diff --check` passed. A new Official Build and installed-wheel verification are still required after merge. This PR does not publish any artifacts. --- pyproject.toml | 1 + tests/unit/packaging/test_wheel_install.py | 31 ++++++++++++++++++++++ 2 files changed, 32 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index 89d7de81b..240dd6c7e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -182,6 +182,7 @@ include = [ "winml", "winml.*" ] "rules/**/*.json", "rules/**/*.parquet", "runtime_checker/need_rerun_errors.json", + "utils/avalizble_ep_device_ops/avaliable_providers.json", ] "winml.modelkit.export" = [ "compatibility_rules.json", diff --git a/tests/unit/packaging/test_wheel_install.py b/tests/unit/packaging/test_wheel_install.py index 79cc5cc1a..86d2eaa69 100644 --- a/tests/unit/packaging/test_wheel_install.py +++ b/tests/unit/packaging/test_wheel_install.py @@ -9,8 +9,10 @@ import shutil import subprocess import sys +import tarfile import textwrap from pathlib import Path +from zipfile import ZipFile import pytest @@ -44,6 +46,35 @@ def _run( return result +def test_built_distributions_include_available_providers(tmp_path: Path) -> None: + """Keep the provider validity mapping in both source and wheel distributions.""" + uv = shutil.which("uv") + assert uv is not None, "uv is required to build the distributions" + config_path = "winml/modelkit/analyze/utils/avalizble_ep_device_ops/avaliable_providers.json" + expected_config = json.loads((REPO_ROOT / "src" / config_path).read_text(encoding="utf-8")) + assert expected_config + dist_dir = tmp_path / "dist" + + _run([uv, "build", "--out-dir", str(dist_dir), str(REPO_ROOT)], cwd=tmp_path) + + wheels = list(dist_dir.glob("winml_cli-*.whl")) + sdists = list(dist_dir.glob("winml_cli-*.tar.gz")) + assert len(wheels) == 1 + assert len(sdists) == 1 + with ZipFile(wheels[0]) as wheel: + assert config_path in wheel.namelist() + assert json.loads(wheel.read(config_path)) == expected_config + with tarfile.open(sdists[0]) as sdist: + config_members = [ + member for member in sdist.getmembers() if member.name.endswith(f"/src/{config_path}") + ] + assert len(config_members) == 1 + config_file = sdist.extractfile(config_members[0]) + assert config_file is not None + with config_file: + assert json.load(config_file) == expected_config + + @pytest.mark.network @pytest.mark.slow @pytest.mark.skipif(sys.platform != "win32", reason="Windows wheel dependencies are required") From efcbdc7a78c1350adde452c8d4fed4ad109aee67 Mon Sep 17 00:00:00 2001 From: "Qiong Wu (qiowu)" Date: Tue, 22 Sep 2026 15:24:33 +0800 Subject: [PATCH 5/6] fix(perf): bind concrete CGC input shapes before compilation (#1433) Local ONNX perf compilation can run before concrete input dimensions reach CGC. Resolve shapes from input metadata or CLI defaults and forward them through Runtime compiler options or ORT named-dimension overrides. Provider-discovery changes are excluded. Windows 11 ARM64 / Adreno X2-90, x64 Python 3.11, WindowsML 2.7.30.dev0 + onnxruntime-windowsml 1.30.0.202609102321. CLI official main 3dd87d3 or PR 1433 at 43d25621. All fixture code is independently authored; no private model, weights, paths or input files are needed. ## Minimal synthetic reproduction Activate the prepared CLI environment with onnx/numpy installed. Save this as make_repro.py: ```python import numpy as np import onnx from onnx import helper as h, TensorProto as T, numpy_helper as nh def tensor(name, dtype, shape): return h.make_tensor_value_info(name, dtype, shape) kind = 'dynamic_add' initializers = [] opset = 18 nodes = [h.make_node('Add', ['a', 'b'], ['y'])] inputs = [tensor('a', T.FLOAT, ['batch', 3]), tensor('b', T.FLOAT, ['batch', 3])] outputs = [tensor('y', T.FLOAT, ['batch', 3])] model = h.make_model(h.make_graph(nodes, kind, inputs, outputs, initializers), opset_imports=[h.make_opsetid('', opset)], ir_version=10) onnx.checker.check_model(model) onnx.save(model, kind + '.onnx') ``` ```powershell python ./make_repro.py winml perf -m ./dynamic_add.onnx --runtime winml-runtime --backend cgc --device gpu --iterations 1 --warmup 0 --no-memory --no-monitor $LASTEXITCODE ``` Also run WinMLCG: ```powershell $env:WINMLCLI_EP_PATH = (Resolve-Path .venv/Lib/site-packages/windowsml/lib).Path winml perf -m ./dynamic_add.onnx --runtime winml-ort --ep winmlcg --device gpu --iterations 1 --warmup 0 --no-memory --no-monitor ``` ## Verified behavior The 117-byte dynamic Add graph uses CLI default batch=1. No shape file or NPZ is required. Base 3dd87d3: A/E both fail compilation. PR 43d25621: A/E both complete using the identical wheels/model. ONNX ReferenceEvaluator validates the elementary computation. This is a shape-handoff test, not full-model support or accuracy certification. Review fixes preserve anonymous input axes for compiler handling, accept matching static zero extents, and read NPY v1/v2/v3 headers without loading array payloads. Named-dimension overrides still require positive int64 values. The Runtime session import is now used explicitly to address the CodeQL unused-import finding. Validation of the review fixes: 197 targeted perf/session/CGC-export tests, repository-wide Ruff, and full-package mypy (464 source files) passed. Regression tests reproduced all three reported failures before the fixes and pass afterward. Native GPU compilation was not rerun for this revision. Only generated data is included. Private inventory references and full-model details have been removed. --- docs/commands/perf.md | 13 ++ src/winml/modelkit/commands/perf.py | 87 ++++++++++ src/winml/modelkit/session/runtime_session.py | 76 ++++++++- .../unit/commands/test_perf_runtime_shapes.py | 158 ++++++++++++++++++ tests/unit/session/test_runtime_session.py | 47 ++++++ 5 files changed, 380 insertions(+), 1 deletion(-) create mode 100644 tests/unit/commands/test_perf_runtime_shapes.py diff --git a/docs/commands/perf.md b/docs/commands/perf.md index 7abe93dd0..bc8045e88 100644 --- a/docs/commands/perf.md +++ b/docs/commands/perf.md @@ -421,6 +421,19 @@ not supported. - **Random inputs do not represent real data distributions.** Latency numbers are accurate, but memory access patterns may differ from production because the generated tensors are uniform random values. For memory-bandwidth-sensitive models this can understate real-world latency. - **Cross-device comparison.** To compare performance across devices, run `winml perf` separately with different `--device` values and compare the resulting JSON reports. +## Concrete input shapes for CGC + +For local ONNX models, Runtime CGC and WinMLCG receive concrete named input +dimensions before compilation. With --input-data, shapes come from NPZ headers +without allocating input tensors. Otherwise, the existing --shape-config and +batch-size resolution rules apply. Rank, static axes, positive dimensions and +shared symbolic names must agree. The source ONNX is not rewritten. + +Runtime CGC requires a Runtime compiler that supports symbolic-dimension +options. Anonymous dynamic axes are rejected rather than guessed, and these +overrides do not resolve internal data-dependent shapes. Input payload loading +and dtype conversion still occur at the normal input-allocation boundary. + ## See also - [winml eval](eval.md) — measure accuracy after benchmarking diff --git a/src/winml/modelkit/commands/perf.py b/src/winml/modelkit/commands/perf.py index 98061bb83..ee4bf2cec 100644 --- a/src/winml/modelkit/commands/perf.py +++ b/src/winml/modelkit/commands/perf.py @@ -907,6 +907,70 @@ def load_input_data( return _load_input_data(path, io_config) +def _runtime_input_shapes( + model_path: Path, input_data: Path | None, shape_config: dict | None, batch_size: int +) -> tuple[dict[str, tuple[int, ...]], dict[str, int]]: + """Resolve source-ONNX shapes before compilation without allocating input tensors.""" + import zipfile + + from ..onnx import get_io_config + from ..session.runtime_session import _symbolic_dimensions_for_inputs + + io_config = get_io_config(model_path) + if not any(dim is None for shape in io_config["input_shapes"] for dim in shape): + return {}, {} + shapes: dict[str, tuple[int, ...]] = {} + if input_data is not None: + if input_data.suffix.lower() != ".npz": + raise click.UsageError("--input-data must be a named .npz archive.") + try: + with zipfile.ZipFile(input_data) as archive: + members = archive.namelist() + expected = [name + ".npy" for name in io_config["input_names"]] + if len(members) != len(expected) or set(members) != set(expected): + raise ValueError("archive keys must exactly match ONNX input names") + for name in io_config["input_names"]: + with archive.open(name + ".npy") as stream: + version = np.lib.format.read_magic(stream) + # NumPy has no public v3 header reader. Use the same + # version-aware, size-limited reader as np.load, without + # reading or allocating the array payload. + shape, _, dtype = np.lib.format._read_array_header( # type: ignore[attr-defined] + stream, version + ) + if dtype.hasobject: + raise ValueError("object input arrays are unsupported") + shapes[name] = shape + except (OSError, ValueError, EOFError, zipfile.BadZipFile) as exc: + raise click.UsageError(f"Cannot read concrete --input-data shapes: {exc}") from exc + else: + for name, shape, symbolic in zip( + io_config["input_names"], + io_config["input_shapes"], + io_config["input_symbolic_shapes"], + strict=True, + ): + full_shape = (shape_config or {}).get(name) + if isinstance(full_shape, (list, tuple)): + # Preserve original values for strict integer/static-axis validation. + shapes[name] = tuple(full_shape) + else: + shapes[name] = _resolve_shape( + shape, name, batch_size, symbolic_shape=symbolic, shape_config=shape_config + ) + return shapes, _symbolic_dimensions_for_inputs(io_config, shapes) + + +def _ort_options_for_dimensions(dimensions: dict[str, int]) -> Any: + """Create fresh ORT options with concrete dimensions before eager session creation.""" + import onnxruntime as ort + + options = ort.SessionOptions() + for name, extent in dimensions.items(): + options.add_free_dimension_override_by_name(name, extent) + return options + + def effective_batch_size( inputs: dict[str, np.ndarray], input_names: list[str], @@ -1361,6 +1425,23 @@ def _load_model(self) -> None: } if is_onnx: + runtime_shapes: dict[str, tuple[int, ...]] = {} + runtime_cgc = self.config.runtime == "winml-runtime" and self._runtime_backend == "cgc" + ort_cgc = ( + self.config.runtime == "winml-ort" + and self._ep_device.device.ep_name == "WinMLCGExecutionProvider" + ) + if runtime_cgc or ort_cgc: + runtime_shapes, dimensions = _runtime_input_shapes( + model_path, + self.config.input_data, + self.config.shape_config, + self.config.batch_size, + ) + if ort_cgc and dimensions: + common_kwargs["session_options"] = lambda: _ort_options_for_dimensions( + dimensions + ) with suppress_native_warnings(enabled=True): self._model = WinMLAutoModel.from_onnx( onnx_path=model_path, @@ -1368,6 +1449,12 @@ def _load_model(self) -> None: compile_provider_options=self.config.compile_ep_options, **common_kwargs, ) + if runtime_cgc and runtime_shapes: + from ..session.runtime_session import WinMLRuntimeSession + + # Keep the lazy import visible to CodeQL as well as type checkers. + runtime_session = cast(WinMLRuntimeSession, self._single._session) # noqa: TC006 + runtime_session.set_input_shapes(runtime_shapes) elif is_mlir: with suppress_native_warnings(enabled=True): self._model = WinMLAutoModel.from_mlir( diff --git a/src/winml/modelkit/session/runtime_session.py b/src/winml/modelkit/session/runtime_session.py index 79fb92a2f..3f19d8db8 100644 --- a/src/winml/modelkit/session/runtime_session.py +++ b/src/winml/modelkit/session/runtime_session.py @@ -30,6 +30,7 @@ import ctypes import json import logging +import operator import threading from contextlib import contextmanager from dataclasses import dataclass @@ -246,6 +247,51 @@ def _stage_schema(wr: Any, stage: Any) -> Any: return schema_type(interface, stage) +def _symbolic_dimensions_for_inputs( + io_config: dict[str, Any], input_shapes: Mapping[str, Any] +) -> dict[str, int]: + """Validate concrete input shapes and resolve shared ONNX dimension names.""" + names = io_config["input_names"] + if set(input_shapes) != set(names): + raise click.ClickException("Concrete input shapes must match the ONNX input names exactly.") + overrides: dict[str, int] = {} + for name, declared, symbolic in zip( + names, io_config["input_shapes"], io_config["input_symbolic_shapes"], strict=True + ): + actual = input_shapes[name] + if len(actual) != len(declared): + raise click.ClickException(f"Input {name!r} rank does not match the ONNX schema.") + for axis, (value, fixed, symbol) in enumerate(zip(actual, declared, symbolic, strict=True)): + try: + extent = operator.index(value) + except TypeError as exc: + raise click.ClickException( + f"Input {name!r} axis {axis} must be an integer." + ) from exc + if isinstance(value, bool) or not 0 <= extent <= (1 << 63) - 1: + raise click.ClickException( + f"Input {name!r} axis {axis} must be a nonnegative int64." + ) + if fixed is not None: + if extent != fixed: + raise click.ClickException( + f"Input {name!r} axis {axis} is {extent}; ONNX requires {fixed}." + ) + elif not isinstance(symbol, str) or not symbol: + # Anonymous axes cannot be bound by name. Let the compiler decide + # whether they need specialization (unused inputs may not). + continue + elif extent == 0: + raise click.ClickException(f"Input {name!r} axis {axis} must be a positive int64.") + elif symbol in overrides and overrides[symbol] != extent: + raise click.ClickException( + f"Conflicting concrete sizes for symbolic dimension {symbol!r}." + ) + else: + overrides[symbol] = extent + return overrides + + 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 @@ -647,12 +693,26 @@ def __init__( self._is_pinned: bool = False self._has_named_bindings = False self._compiled_artifacts: TemporaryDirectory[str] | None = None + self._symbolic_dimensions: dict[str, int] = {} self._built = False # Perf tracking, enabled inside perf(). self._perf_stats: PerfStats | None = None # -- lifecycle ---------------------------------------------------------- + def set_input_shapes(self, input_shapes: Mapping[str, Any]) -> None: + """Set validated ONNX dimensions before compiling; never rewrite the model.""" + from ..onnx import get_io_config + + with self._lock: + if self._built: + raise ValueError("Input shapes must be configured before Runtime compilation.") + if self._is_mlir or self._backend != "cgc": + raise ValueError("Concrete compilation shapes require source ONNX and backend=cgc.") + self._symbolic_dimensions = _symbolic_dimensions_for_inputs( + get_io_config(self._model_path), input_shapes + ) + def _ensure_built(self) -> None: """Import the runtime, resolve the target, load and build the pipeline. @@ -695,9 +755,23 @@ def _load_onnx_on_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() + options = None try: - compiler.compile_to_file(source_model, str(artifact_path)) + if self._symbolic_dimensions: + options = compiler.create_options() + if not options.supports_symbolic_dimensions: + raise click.ClickException( + "The installed Runtime compiler does not support " + "symbolic dimensions." + ) + for name, extent in self._symbolic_dimensions.items(): + options.symbolic_dimensions[name] = extent + compiler.compile_to_file(source_model, str(artifact_path), options=options) + else: + compiler.compile_to_file(source_model, str(artifact_path)) finally: + if options is not None: + options.close() compiler.close() with _translate_native_errors("load"): diff --git a/tests/unit/commands/test_perf_runtime_shapes.py b/tests/unit/commands/test_perf_runtime_shapes.py new file mode 100644 index 000000000..059479b59 --- /dev/null +++ b/tests/unit/commands/test_perf_runtime_shapes.py @@ -0,0 +1,158 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Concrete compilation shape handoff, with no native execution.""" + +from io import BytesIO +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import Mock +from zipfile import ZipFile + +import click +import numpy as np +import onnx +import pytest + +from winml.modelkit.commands.perf import ( + BenchmarkConfig, + PerfBenchmark, + _ort_options_for_dimensions, + _runtime_input_shapes, +) + + +def _model(tmp_path: Path) -> Path: + inputs = [ + onnx.helper.make_tensor_value_info(name, onnx.TensorProto.FLOAT, ["batch", 3]) + for name in ("left", "right") + ] + output = onnx.helper.make_tensor_value_info("out", onnx.TensorProto.FLOAT, ["batch", 3]) + model = onnx.helper.make_model( + onnx.helper.make_graph( + [onnx.helper.make_node("Add", ["left", "right"], ["out"])], "shapes", inputs, [output] + ) + ) + path = tmp_path / "model.onnx" + onnx.save(model, path) + return path + + +def test_npz_headers_override_shape_config_without_loading_tensors(tmp_path, monkeypatch): + model = _model(tmp_path) + before = model.read_bytes() + data = tmp_path / "inputs.npz" + np.savez(data, left=np.zeros((2, 3)), right=np.zeros((2, 3))) + monkeypatch.setattr(np, "load", Mock(side_effect=AssertionError("No tensor allocation"))) + shapes, dimensions = _runtime_input_shapes(model, data, {"batch": 99}, 42) + assert shapes == {"left": (2, 3), "right": (2, 3)} + assert dimensions == {"batch": 2} + assert model.read_bytes() == before + + +def test_shape_config_resolves_symbols_without_inputs(tmp_path): + shapes, dimensions = _runtime_input_shapes(_model(tmp_path), None, {"batch": 4}, 1) + assert shapes == {"left": (4, 3), "right": (4, 3)} + assert dimensions == {"batch": 4} + + +@pytest.mark.parametrize( + "unused_shape, actual_shape", [([None, 3], (2, 3)), (["extent", 0], (2, 0))] +) +def test_unused_input_preserves_anonymous_and_static_zero_axes( + tmp_path, unused_shape, actual_shape +): + path = _model(tmp_path) + model = onnx.load(path) + model.graph.input.append( + onnx.helper.make_tensor_value_info("unused", onnx.TensorProto.FLOAT, unused_shape) + ) + onnx.save(model, path) + data = tmp_path / "inputs.npz" + np.savez(data, left=np.zeros((2, 3)), right=np.zeros((2, 3)), unused=np.zeros(actual_shape)) + shapes, dimensions = _runtime_input_shapes(path, data, None, 1) + assert shapes["unused"] == actual_shape + assert dimensions == ({"batch": 2, "extent": 2} if actual_shape[1] == 0 else {"batch": 2}) + + +@pytest.mark.parametrize("version", [(1, 0), (2, 0), (3, 0)]) +def test_npz_header_versions_without_reading_payload(tmp_path, version): + data = tmp_path / "headers.npz" + with ZipFile(data, "w") as archive: + for name in ("left", "right"): + stream = BytesIO() + np.lib.format.write_array(stream, np.zeros((2, 3), dtype=np.float32), version=version) + # Omit the payload: shape discovery must only read the header. + archive.writestr(name + ".npy", stream.getvalue()[:-24]) + shapes, dimensions = _runtime_input_shapes(_model(tmp_path), data, None, 1) + assert shapes == {"left": (2, 3), "right": (2, 3)} + assert dimensions == {"batch": 2} + + +@pytest.mark.parametrize("right_shape", [(3, 3), (2, 4), (2, 3, 1)]) +def test_npz_rejects_symbol_conflict_static_mismatch_and_rank(tmp_path, right_shape): + data = tmp_path / "inputs.npz" + np.savez(data, left=np.zeros((2, 3)), right=np.zeros(right_shape)) + with pytest.raises(click.ClickException): + _runtime_input_shapes(_model(tmp_path), data, None, 1) + + +def test_ort_options_are_fresh_and_receive_dimensions_before_use(monkeypatch): + import onnxruntime as ort + + created = [] + + def create(): + value = SimpleNamespace(add_free_dimension_override_by_name=Mock()) + created.append(value) + return value + + monkeypatch.setattr(ort, "SessionOptions", create) + first = _ort_options_for_dimensions({"batch": 2}) + second = _ort_options_for_dimensions({"batch": 2}) + assert first is not second + for options in created: + options.add_free_dimension_override_by_name.assert_called_once_with("batch", 2) + + +@pytest.mark.parametrize("runtime", ["winml-runtime", "winml-ort"]) +def test_perf_hands_shapes_to_runtime_before_compilation(tmp_path, monkeypatch, runtime): + from winml.modelkit.models import WinMLAutoModel + + model_path = _model(tmp_path) + data = tmp_path / "inputs.npz" + np.savez(data, left=np.zeros((2, 3)), right=np.zeros((2, 3))) + config = BenchmarkConfig( + model_id=str(model_path), + runtime=runtime, + backend="cgc" if runtime == "winml-runtime" else None, + input_data=data, + device="gpu", + ) + benchmark = PerfBenchmark(config) + benchmark._ep_device = SimpleNamespace( + device=SimpleNamespace(ep_name="WinMLCGExecutionProvider") + ) + monkeypatch.setattr(benchmark, "_resolve_device_ep", lambda: None) + session = SimpleNamespace(set_input_shapes=Mock(), compile=Mock()) + wrapper = SimpleNamespace(_session=session) + options = object() + configure = Mock(return_value=options) + monkeypatch.setattr("winml.modelkit.commands.perf._ort_options_for_dimensions", configure) + + def construct(**kwargs): + if runtime == "winml-ort": + assert kwargs["session_options"]() is options + configure.assert_called_once_with({"batch": 2}) + else: + assert "session_options" not in kwargs + return wrapper + + monkeypatch.setattr(WinMLAutoModel, "from_onnx", construct) + benchmark._load_model() + if runtime == "winml-runtime": + session.set_input_shapes.assert_called_once_with({"left": (2, 3), "right": (2, 3)}) + else: + session.set_input_shapes.assert_not_called() + session.compile.assert_not_called() diff --git a/tests/unit/session/test_runtime_session.py b/tests/unit/session/test_runtime_session.py index 897cac4ec..f5ad8b673 100644 --- a/tests/unit/session/test_runtime_session.py +++ b/tests/unit/session/test_runtime_session.py @@ -24,6 +24,7 @@ _numpy_dtype_for, _ResolvedRuntimeTarget, _shape_with_dynamic_dims, + _symbolic_dimensions_for_inputs, ) @@ -206,6 +207,52 @@ def load_model(path: str) -> Any: session.reset() +@pytest.mark.parametrize("fails", [False, True]) +def test_cgc_compile_uses_public_symbolic_options_and_closes_them(monkeypatch, fails): + io = { + "input_names": ["x"], + "input_shapes": [[None, 3]], + "input_symbolic_shapes": [["batch", 3]], + } + monkeypatch.setattr("winml.modelkit.onnx.get_io_config", lambda _path: io) + options = SimpleNamespace( + supports_symbolic_dimensions=True, symbolic_dimensions={}, close=Mock() + ) + source = Mock() + compiler = Mock(create_options=Mock(return_value=options)) + if fails: + compiler.compile_to_file.side_effect = RuntimeError("compile failed") + runtime = Mock(load_model=Mock(side_effect=[source, _Model()])) + target = _ResolvedRuntimeTarget( + execution_target=SimpleNamespace(model_compiler=lambda: compiler), device_class="gpu" + ) + session = WinMLRuntimeSession("source.onnx", ep_device=_mlir_ep_device(), backend="cgc") + session.set_input_shapes({"x": (2, 3)}) + try: + if fails: + with pytest.raises(RuntimeError, match="compile failed"): + session._load_onnx_on_cgc(runtime, target) + else: + session._load_onnx_on_cgc(runtime, target) + assert options.symbolic_dimensions == {"batch": 2} + assert compiler.compile_to_file.call_args.kwargs == {"options": options} + options.close.assert_called_once_with() + compiler.close.assert_called_once_with() + finally: + session.reset() + + +@pytest.mark.parametrize("shape", [(True, 3), (1.5, 3), (0, 3), (1 << 63, 3)]) +def test_symbolic_dimensions_reject_invalid_extents(shape): + io = { + "input_names": ["x"], + "input_shapes": [[None, 3]], + "input_symbolic_shapes": [["batch", 3]], + } + with pytest.raises(click.ClickException): + _symbolic_dimensions_for_inputs(io, {"x": shape}) + + def test_schema_helpers() -> None: assert _numpy_dtype_for(SimpleNamespace(name="FLOAT16")) == "float16" assert _shape_with_dynamic_dims([0, -1, (1 << 64) - 1, 4]) == [ From b8d9073dc0c55307a9eaaffaa8a6271e0d75dad3 Mon Sep 17 00:00:00 2001 From: Zac <1221537+tezheng@users.noreply.github.com> Date: Wed, 23 Sep 2026 17:24:20 +0800 Subject: [PATCH 6/6] =?UTF-8?q?feat(cgc):=20inspect=20D3D12=20adapters=20a?= =?UTF-8?q?nd=20the=20MLIR=20patterns=20their=20drivers=E2=80=A6=20(#1454)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit … declare Two questions get asked on every new machine and after every driver update: which D3D12 adapters are here, and what does this driver's MLIR-program implementation claim it can match. Both are answered through public Windows APIs with ctypes, so nothing has to be built and no extra package installed. winml cgc adapters # which adapters are here winml cgc patterns [-a ADAPTER] [-v] # what a driver declares winml cgc patterns -a ADAPTER --dump [--overwrite] # keep it on disk winml cgc patterns [-a ADAPTER] --open [FILE] # look at a dump Listing adapters needs nothing beyond Windows. Anything touching MLIR patterns needs an Agility SDK D3D12Core.dll of SDK 720 or newer, named by --d3d12-dir or $WINML_D3D12_DIR/$D3D12_DIR, or dropped in bin/ beside the venv or the install; the inbox runtime does not serve the preview feature. Which core answered opens every run on the redist: line, because the same driver reports different things through the 720 and 721 exchange shapes. When none can be found, the MLIR column reads ? rather than no -- the question could not be asked -- and the failure names every place that was searched. With no -a, `patterns` asks every adapter and prints one line each; a driver that claims the exchange and then refuses it costs only its own line, on stderr, and the run exits 1 at the end. A -a number is an index when one exists and otherwise a description substring, so `-a 5060` finds the RTX 5060. A dump is filed under patterns/// and holds patterns.mlir, metadata.txt (the six keys dxcgc-dump-driver-patterns writes, plus which redist and ABI produced it and how the stored payload differs from the one received) and patterns.json: one record per pattern -- index, name, kind, benefit, rule count, source file, kernel, line count and its own text -- plus a summary, per-kind and per-source grouping and a format tag, so nothing downstream has to parse MLIR. Text payloads are normalised once, where the exchange returns them: the trailing NUL is dropped and every line break becomes LF, so the counts, patterns.mlir and patterns.json describe the same bytes. Bytecode is left alone. Counting treats a pattern as one cgc_pattern.pattern declaration and a rule as one flat match alternative after any_of expansion, because comparing declaration counts alone across drivers can read a loss where the matching surface grew. Comments are blanked rather than deleted so offsets stay valid, the displayed text keeps its comments, and a brace inside a comment cannot unbalance a slice. --open shows a dump in the CGC Pattern Atlas, a self-contained page that ships inside the package and is resolved with importlib.resources, so it works from a wheel, an editable install and a checkout alike. A dump is embedded into a copy of the page named after it, with every "<" escaped so a pattern's own text cannot close the script element. --open FILE shows a dump taken earlier and asks no driver anything; anything that is not a patterns.json is a usage error. Verified on an RTX 5060 (driver 32.0.16.3004), an AMD 880M and WARP against the SDK 720 and 721 redists in DXCGC-Drop10.0-x64 and against Drop11: 51 patterns / 63 rules on the 720 path, byte-identical dumps across refactors, and the documented exit codes on every failure path. docs/commands/cgc.md covers both commands, every flag and scenario, the dump layout and patterns.json. TODO(tests) markers record the gaps left untested. --- .gitignore | 9 + docs/commands/cgc.md | 682 ++++++ docs/commands/overview.md | 5 +- mkdocs.yml | 1 + pyproject.toml | 3 + .../commands/assets/cgc-pattern-atlas.html | 608 +++++ src/winml/modelkit/commands/cgc.py | 1986 +++++++++++++++++ tests/unit/commands/test_cgc.py | 1299 +++++++++++ 8 files changed, 4592 insertions(+), 1 deletion(-) create mode 100644 docs/commands/cgc.md create mode 100644 src/winml/modelkit/commands/assets/cgc-pattern-atlas.html create mode 100644 src/winml/modelkit/commands/cgc.py create mode 100644 tests/unit/commands/test_cgc.py diff --git a/.gitignore b/.gitignore index 8413bf4d0..fd4910374 100644 --- a/.gitignore +++ b/.gitignore @@ -281,3 +281,12 @@ src/winml/modelkit/analyze/rules/runtime_check_rules/**/*.parquet # Generated by mike (docs versioning) docs/versions.json /uv.lock + +# An Agility SDK redist dropped where `winml cgc` looks for one, beside the venv +# or beside the installed package (see docs/commands/cgc.md). Multi-MB +# D3D12Core.dll: placed locally, never committed. +/bin/ +/src/winml/bin/ + +# `winml cgc patterns --dump` writes into the working directory. +/patterns/ diff --git a/docs/commands/cgc.md b/docs/commands/cgc.md new file mode 100644 index 000000000..ccd5397f2 --- /dev/null +++ b/docs/commands/cgc.md @@ -0,0 +1,682 @@ +# winml cgc + +> Which D3D12 adapters are on this machine, and what patterns does an adapter's driver +> declare it can match. + +## When to use this + +Two questions get asked on every new machine and after every driver update: which +adapters are here, and what does this driver's MLIR-program implementation claim it can +match. `winml cgc adapters` answers the first and `winml cgc patterns` the second, both +from public Windows APIs through `ctypes`, so nothing has to be built. + +Use it to check whether a GPU's driver implements D3D MLIR programs before relying on +the CGC path, to record what a driver declares after a driver update, and to compare +two drivers' matching surfaces. + +## Synopsis + +```bash +$ winml cgc adapters [options] # which adapters are here +$ winml cgc patterns [options] # what a driver declares it can match +``` + +`cgc` is a command group; run `winml cgc` or `winml cgc --help` to list its subcommands. +Both print help and exit 0. + +## Requirements + +Listing adapters needs nothing beyond Windows. + +Querying a driver's MLIR support or patterns needs an **Agility SDK `D3D12Core.dll` of +SDK 720 or newer**, matching the Python interpreter's architecture, and a compatible +driver. The MLIR-program feature is a preview feature that the inbox D3D12 runtime does +not serve. See [Finding the redist](#finding-the-redist). + +`winml cgc patterns --open` without `--dump` only opens the viewer or an existing JSON +dump. It needs a browser, but neither a GPU nor a redist, and does not query drivers. + +## Flags + +Shared by both subcommands: + +| Flag | Short | Type | Default | Description | +|------|-------|------|---------|-------------| +| `--d3d12-dir` | | path | `""` | Directory holding `D3D12Core.dll`. Overrides every other source. | +| `--verbose` | `-v` | count | `0` | Increase logging verbosity (`-v` for INFO, `-vv` for DEBUG). Any nonzero level also adds raw capability HRESULTs and, under `patterns`, a per-pattern listing. Accepted at the root or subcommand, so `winml -v cgc adapters` and `winml cgc adapters -v` behave the same. | +| `--quiet` | `-q` | flag | `false` | Errors only on the logging channel. Does not silence stdout or the directly printed probe details requested by `-v`. | +| `--help` | `-h` | flag | — | Show help and exit. | + +`winml cgc patterns` adds: + +| Flag | Short | Type | Default | Description | +|------|-------|------|---------|-------------| +| `--adapter` | `-a` | string | `""` | The adapter to ask, by description substring (case-insensitive) or by index. A number that is not an index is tried as a substring, so `-a 5060` names an RTX 5060. Without it, every enumerated adapter is asked. Required by `--dump`. | +| `--dump` | | flag | `false` | Write a dump under `./patterns`. The files depend on the response: text, bytecode, empty, or unsupported; see [Dump](#dump). Needs `-a`. | +| `--overwrite/--no-overwrite` | | flag | `false` | Replace an existing dump directory. | +| `--open` | | flag or path | — | Open the [pattern atlas](#browsing-a-dump) in a browser. With `--dump`, show this run's dump when available. Given a `patterns.json`, show that file. With neither a file nor `--dump`, open an empty atlas. Without `--dump`, no driver is queried, so neither a redist nor a GPU is needed. | + +## Usage scenarios + +Common tasks, including viewing a dump without querying a driver. + +| You want to know | Run | +|---|---| +| which adapters are on this machine, and whose driver implements MLIR programs | `winml cgc adapters` | +| *why* an adapter is marked `no` | `winml cgc adapters -v` | +| which Agility core a live query uses | the `redist:` line of a run that reaches redist resolution; viewer-only runs do not resolve one | +| how many patterns and rules this driver declares | `winml cgc patterns -a nvidia` | +| what every enumerated adapter's driver declares | `winml cgc patterns` | +| every pattern it declares, and what kind of claim each one is | `winml cgc patterns -a nvidia -v` | +| whether a non-NVIDIA adapter claims anything | `winml cgc patterns -a amd` | +| keep the declaration on disk, for a diff or a record | `winml cgc patterns -a nvidia --dump` | +| re-capture after a driver update | `winml cgc patterns -a nvidia --dump --overwrite` | +| dump a driver's patterns and browse them in the atlas | `winml cgc patterns -a nvidia --dump --open` | +| browse an existing dump without querying drivers | `winml cgc patterns --open ` | +| open an empty atlas to load a dump manually | `winml cgc patterns --open` | +| what a *different* redist makes the same driver say | `winml cgc patterns -a nvidia --d3d12-dir ` | +| record that an adapter has *no* MLIR support, for the same baseline | `winml cgc patterns -a amd --dump` | +| which subcommands the group has | `winml cgc --help` | + +The output examples below illustrate the format. Adapter names, versions, counts, and +HRESULTs depend on the machine, driver, and selected redist; they are not compatibility +guarantees. + +### A machine with no redist + +Listing adapters needs no redist, so this is the first thing to run on a new box. +Without a redist the capability question cannot be asked, and the column says so rather +than guessing: + +```console +$ winml cgc adapters +redist: none -- MLIR support unknown (run with -v to see where cgc looked) +IDX ADAPTER DRIVER TYPE ATTRIBUTES MLIR +-------------------------------------------------------------------------------- +0 NVIDIA GeForce RTX 5060 Laptop 32.0.16.3004 hardware ML CC GFX ? +1 AMD Radeon(TM) 880M Graphics 32.0.31041.1004 integrated ML CC GFX ? +2 Microsoft Basic Render Driver 10.0.26100.9278 software ML CC GFX ? +``` + +`?` is not `no`. The driver versions and the table are still correct and still useful. + +### List adapters + +The command enumerates DXCore adapters advertising generic ML support. If that list +is empty, it retries with core-compute support; it does not combine the two lists or +enumerate every graphics-only adapter. The list is sorted hardware-first, then by +high-performance preference. + +```console +$ winml cgc adapters --d3d12-dir C:\path\to\D3D12 +redist: C:\path\to\D3D12 (SDK 720, x64) +IDX ADAPTER DRIVER TYPE ATTRIBUTES MLIR +-------------------------------------------------------------------------------- +0 NVIDIA GeForce RTX 5060 Laptop 32.0.16.3004 hardware ML CC GFX yes +1 AMD Radeon(TM) 880M Graphics 32.0.31041.1004 integrated ML CC GFX no +2 Microsoft Basic Render Driver 10.0.26100.9278 software ML CC GFX no +``` + +| Column | Meaning | +|---|---| +| `IDX` | Selection index accepted by `patterns -a`. | +| `ADAPTER` | DXCore description, truncated to 30 characters in this table only. | +| `DRIVER` | Driver version as four 16-bit components. | +| `TYPE` | `hardware`, `integrated`, or `software`. | +| `ATTRIBUTES` | `ML`: generic ML; `CC`: core compute; `GFX`: D3D12 graphics. These are not MLIR support indicators. | +| `MLIR` | `yes` or `no` from the capability probe, or `?` when it could not be asked through a usable redist. | + +The `redist:` line precedes the table when redist resolution is reached. An explicitly +named directory rejected by the on-disk checks fails before that line; a DLL that +passes those checks but is rejected by D3D12 fails afterwards. See +[Driving it from a script](#driving-it-from-a-script) for output on failed runs. + +`MLIR` is the column that matters: it is what the driver claims about the MLIR-program +exchange through the redist on the `redist:` line. A driver may report support through +a 721-shaped query but refuse the subsequent pattern exchange; see +[When something looks wrong](#when-something-looks-wrong). Fetching a declaration with +`patterns` checks that exchange, not whether a particular model will compile. + +Probing support attempts to create a D3D12 device for each adapter and can take longer +than enumerating adapter metadata. Its cost depends on the machine and driver. + +### Counts + +```console +$ winml cgc patterns -a nvidia +redist: C:\path\to\D3D12 (SDK 720, x64) +NVIDIA GeForce RTX 5060 Laptop GPU: 51 patterns / 63 rules +``` + +`-a` answers for **one** adapter, so it reports that adapter rather +than printing the table. Add `-v` to list every pattern the driver declares, grouped by +what kind of claim it is: + +```console +$ winml cgc patterns -a nvidia -v +redist: C:\path\to\D3D12 (SDK 720, x64) +NVIDIA GeForce RTX 5060 Laptop GPU: 51 patterns / 63 rules + +kernel 12 patterns, 16 rules declares a jitFunction -- the only kind that claims work + 38 gemm_input_major_expression_epilogue_cluster + 36 gemm_input_major 3 rules + ... +fusion 4 patterns, 6 rules collapses a marked cluster into one native op + 200 lora_no_views_output_major 3 rules + ... +hint 5 patterns, 9 rules anchors a cluster and sets the order clusters grow in +mark 25 patterns, 27 rules tags an op so a fusion root can absorb it later +rewrite 5 patterns, 5 rules graph surgery, so that other patterns can match +``` + +The five kinds are read off each pattern's body, never its name: a `jitFunction` makes +it a **kernel**, a `subgraph_rewrite_desc` without one a **fusion**, then +`cgc_add_pattern_cluster_hint` a **hint**, `cgc_mark_pattern_cluster_op` a **mark**, and +anything left a **rewrite**. Groups run from most claimed to least, patterns within a +group by descending benefit, then by full name for ties. A missing benefit is sorted +as zero and displayed as `-`. A per-pattern rule count appears only where a pattern +expands to more than one rule. A kind the driver declares nothing for is left out. + +On an adapter whose driver has no MLIR support, this is a normal answer, not an error +(**exit 0**): + +```console +$ winml cgc patterns -a amd +redist: C:\path\to\D3D12 (SDK 720, x64) +AMD Radeon(TM) 880M Graphics: driver does not implement MLIR programs +``` + +Without a redist the question cannot be asked at all, and that is a failure rather +than an answer (**exit 1**): + +```console +$ winml cgc patterns -a nvidia +redist: none -- MLIR support unknown (run with -v to see where cgc looked) +Error: cgc patterns needs an Agility SDK redist; none was found (run with -v to see where cgc looked) +``` + +### Dump + +`--dump` must name its adapter with `-a`: + +```console +$ winml cgc patterns -a nvidia --dump +redist: C:\path\to\D3D12 (SDK 720, x64) +NVIDIA GeForce RTX 5060 Laptop GPU: 51 patterns / 63 rules +wrote patterns\nvidia-geforce-rtx-5060-laptop-gpu\32.0.16.3004\patterns.mlir (124753 bytes) +wrote patterns\nvidia-geforce-rtx-5060-laptop-gpu\32.0.16.3004\patterns.json +wrote patterns\nvidia-geforce-rtx-5060-laptop-gpu\32.0.16.3004\metadata.txt +``` + +For a non-empty text response, three files are written: + +| file | what it is | +|---|---| +| `patterns.mlir` | the driver's declaration as text, with LF line endings and without the C string terminator the driver appends | +| `patterns.json` | the same patterns split one per record and grouped — see [patterns.json](#patternsjson) | +| `metadata.txt` | provenance: adapter, driver, redist, SDK, and the counts | + +Other responses produce different files: + +| Response | Files written by `--dump` | +|---|---| +| MLIR bytecode | `patterns.mlirbc` and `metadata.txt`; no counts or `patterns.json`. | +| Empty declaration after text normalization | `metadata.txt` with `status=empty`. | +| No MLIR support reported by the probe | `metadata.txt` with `status=unsupported`. | +| Redist or pattern exchange failure | No new dump is written; the command exits 1. | + +Naming the adapter is required rather than defaulted, because a dump is filed under a +slug and a driver version: dumping whichever adapter happened to be first writes a +correct-looking directory for the wrong device. Omitting `-a` is a usage error: + +```console +$ winml cgc patterns --dump +Usage: winml cgc patterns [OPTIONS] +Try 'winml cgc patterns --help' for help. + +Error: --dump needs an adapter: name one with -a # exit 2 +``` + +The layout is `patterns///`, relative to the working +directory. It is keyed by the normalized description and driver version, not by a +unique device ID or SDK version. Identical descriptions, descriptions that normalize +to the same slug, and different redists used with the same driver can therefore share +a destination. Use separate working directories to keep those captures independently. +The root is always `patterns/`; there is no output-directory flag. + +A dump is **never silently overwritten**: + +```console +$ winml cgc patterns -a nvidia --dump # second time +Error: Dump directory 'patterns\nvidia-geforce-rtx-5060-laptop-gpu\32.0.16.3004' already exists and is not empty. Re-run with --overwrite to replace its contents. +``` + +That message goes to stderr and the exit status is 1. After a successful query, +`--overwrite` removes the four known dump files (`patterns.mlir`, `patterns.mlirbc`, +`patterns.json`, and `metadata.txt`) before writing the replacement; unrelated files +are preserved. A query failure leaves the previous dump untouched, but a later +filesystem write failure can still leave an incomplete replacement. A non-empty +directory or an existing file at the destination blocks without `--overwrite`; an +empty directory is reused. + +`metadata.txt` alongside the dump: + +```ini +adapter=NVIDIA GeForce RTX 5060 Laptop GPU +driver_version=32.0.16.3004 +status=dumped +received_encoding=text +text_file=patterns.mlir +bytecode_file= +abi=720 +sdk_version=720 +redist=C:\path\to\D3D12 +patterns=51 +rules=63 +received_bytes=127117 +line_endings=lf +``` + +The first six keys are the ones `dxcgc-dump-driver-patterns.exe` writes, unchanged so +existing readers keep working; the rest record which redist and ABI produced the dump. + +| key | values | meaning | +|---|---|---| +| `status` | `dumped` \| `unsupported` \| `empty` | whether patterns were captured, the driver has no MLIR support, or it returned an empty declaration | +| `received_encoding` | `text` \| `bytecode` \| `none` | the stored payload's encoding; `none` for `unsupported` or `empty` | +| `text_file` / `bytecode_file` | filename or empty | only one is ever set | +| `abi` / `sdk_version` | e.g. `720` | both currently record the SDK version; SDK 720 uses the 720 exchange shape, SDK 721 or newer uses the 721 shape | +| `patterns` / `rules` | integers, or empty | empty when the counts could not be taken (`unsupported`, `empty`, or a bytecode answer) | +| `received_bytes` | integer, or empty | size returned by the exchange, including any terminator; an empty response records `0` (or the terminator's size), while an unsupported adapter leaves this empty | +| `line_endings` | `lf`, or empty | `lf` when a text payload was normalised to LF line endings and its terminator removed before it was stored | + +#### Dumping an adapter with no MLIR support + +This is a normal answer, not an error (**exit 0**). No `patterns.mlir` is written — just +a `metadata.txt` recording that this driver was asked and said no, so the baseline covers +every adapter rather than silently omitting the ones that declined: + +```console +$ winml cgc patterns -a amd --dump +redist: C:\path\to\D3D12 (SDK 720, x64) +AMD Radeon(TM) 880M Graphics: driver does not implement MLIR programs +wrote patterns\amd-radeon-880m-graphics\32.0.31041.1004\metadata.txt +``` + +```ini +adapter=AMD Radeon(TM) 880M Graphics +driver_version=32.0.31041.1004 +status=unsupported +received_encoding=none +text_file= +bytecode_file= +abi=720 +sdk_version=720 +redist=C:\path\to\D3D12 +patterns= +rules= +received_bytes= +line_endings= +``` + +#### The adapter slug + +The directory name is the description lowercased, with `(R)`, `(TM)` and `(C)` removed +and every run outside ASCII `a-z` and `0-9` collapsed to a single `-`, then any leading +or trailing `-` trimmed. If nothing remains, the slug is `adapter`: + +| description | slug | +|---|---| +| `NVIDIA GeForce RTX 5060 Laptop GPU` | `nvidia-geforce-rtx-5060-laptop-gpu` | +| `AMD Radeon(TM) 880M Graphics` | `amd-radeon-880m-graphics` | +| `Intel(R) Graphics` | `intel-graphics` | + +### After a driver update + +Three things move: the version in the table, the counts, and the dump. In that order: + +```bash +winml cgc adapters # current driver versions +winml cgc patterns -a nvidia # current pattern and rule counts +winml cgc patterns -a nvidia -v # inspect the current declarations +winml cgc patterns -a nvidia --dump # keep the current declaration +``` + +The dump lands under the new driver version, so the previous one is still there beside +it. Compare the two with `git diff --no-index OLD NEW`. Re-dumping the *same* version +is refused until you pass `--overwrite`. + +Text dumps are stored with LF line endings and no trailing NULs; bytecode is preserved +unchanged. Older text dumps may retain CRLF endings and a trailing NUL. Compare +normalized copies: remove trailing NUL bytes and normalize line endings first. +`git diff --no-index --ignore-cr-at-eol OLD NEW` only ignores end-of-line CRs; it does +not remove a trailing NUL, which can make Git treat an older text dump as binary. + +### Selecting an adapter + +`-a` takes a description substring (case-insensitive) or a bare index: + +```bash +winml cgc patterns -a nvidia +winml cgc patterns -a 0 +``` + +With no `-a`, `winml cgc patterns` asks **every enumerated adapter** and prints one line +each. A driver that cannot answer -- one that claims the +exchange and then refuses it -- costs only its own line, on stderr; the rest are still +reported and the run then exits 1. `--dump` does not accept that default: it names +its adapter with `-a` or it exits 2. + +A `-a` value is matched as an index when one exists, and otherwise as a description +substring, so `-a 5060` finds an RTX 5060 on a three-adapter machine. If several +descriptions match, the first in enumeration order wins; use `IDX` to distinguish +them. Matching uses the full description, not the table's truncated display. + +Once adapters and a usable redist are available, a selector that matches nothing +exits 2. `winml cgc adapters` takes no `-a`: it always prints the enumerated list. +If enumeration finds no adapters after argument checks, either live-query subcommand prints +`no D3D12 adapters found` and exits 0 without resolving a redist. + +### Driving it from a script + +`stdout` carries the `redist:` line, table, counts, verbose pattern listing, and +`wrote`/`opening` notices. `stderr` carries logging and diagnostics such as the `-v` +probes and redist failure report. This is human-readable output, not a JSON stream: + +```powershell +winml cgc patterns -a nvidia > counts.txt +if ($LASTEXITCODE -ne 0) { "command failed; output may be partial" } +``` + +The `redist:` line is written after redist resolution succeeds or concludes that no +implicit candidate is available. It remains in stdout if a later step fails. Parsing +errors and an explicitly named redist rejected during discovery fail before that +line. Viewer-only runs do not print it. + +**A failed multi-adapter run can contain partial results.** Successful adapters are +still printed when another adapter's exchange fails, and the command exits 1 after +the sweep. A selected-adapter failure can leave just the `redist:` line. Always check +the exit code; neither non-empty output nor a count line proves the whole run +succeeded. + +Read `patterns=` and `rules=` from a successfully written dump's `metadata.txt`, or +the structured `patterns.json`, rather than re-parsing stdout. There is no JSON-output +flag for these subcommands. + +## patterns.json + +A text dump also writes `patterns.json`: every pattern in the declaration, split out +one per record and grouped, so a reader never has to parse MLIR. The +[CGC Pattern Atlas](#browsing-a-dump) page is built on it. + +```json +{ + "format": "winml-cgc-patterns/1", + "meta": { + "adapter": "NVIDIA GeForce RTX 5060 Laptop GPU", + "driver_version": "32.0.16.3004", + "sdk_version": 720, + "redist": "C:\\path\\to\\D3D12", + "encoding": "text", + "bytes": 124753 + }, + "summary": { + "patterns": 51, "rules": 63, "kernel_patterns": 12, + "distinct_kernels": 9, "sources": 4, "top_benefit": 200 + }, + "kinds": { "kernel": 12, "fusion": 4, "hint": 5, "mark": 25, "rewrite": 5 }, + "sources": [ + { "name": "nvidia_conv_central.pdll", "patterns": [1, 2, 3] } + ], + "patterns": [ + { + "index": 28, + "name": "nvidia.gemm_input_major_expression_epilogue_cluster", + "short_name": "gemm_input_major_expression_epilogue_cluster", + "kind": "kernel", + "benefit": 38, + "rules": 1, + "source": "nvidia_handwritten.mlir", + "kernel": "GemmCluster", + "lines": 92, + "mlir": "cgc_pattern.pattern @nvidia.gemm_input_major_expression_epilogue_cluster ..." + } + ] +} +``` + +(`sources` and `patterns` are abbreviated here.) + +| field | derived from | +|---|---| +| `meta.bytes` | length of the normalized text payload, not the original `received_bytes` in `metadata.txt` | +| `index`, `short_name` | one-based declaration order, and the portion of `name` after its final `.` | +| `kind`, `benefit`, `rules` | the same parser that prints the counts; `benefit` is `null` if absent | +| `source` | the preceding `// from .pdll` / `// from .mlir` marker, or `null` if none exists; prose beginning with "from" does not open a section | +| `kernel` | the value matched from `jitFunction = #cgc.string<"Name">` or a plain string; `null` if no kernel name is parsed | +| `lines`, `mlir` | line count and text from the declaration through its closing brace, including comments; invalid UTF-8 is decoded with replacement characters | +| `summary.kernel_patterns`, `summary.distinct_kernels` | records with a parsed kernel name, and the number of distinct names | +| `sources`, `summary.sources` | groups of pattern indices by source; records without a source share the `(unattributed)` group, which is included in the count | + +`format` names the layout, so a reader can refuse a file it does not understand. A +bytecode dump gets no `patterns.json` — splitting it needs the real MLIR parser — and +neither does an unsupported adapter or an empty declaration. Non-empty text with no +recognized pattern declarations does produce JSON with an empty `patterns` array. + +### Browsing a dump + +The **CGC Pattern Atlas** is a single self-contained HTML page, opened by `--open`: + +```bash +winml cgc patterns -a nvidia --dump --open # dump this driver, then show it +winml cgc patterns --open patterns.json # show a dump taken earlier +winml cgc patterns --open # open an empty page; no driver query +``` + +Without `--dump`, `--open` bypasses adapter and redist selection. Use the file form +without `--dump`: when `--dump` is present, the viewer uses this run's output rather +than any filename passed to `--open`. If no `patterns.json` was written (bytecode, +empty declaration, or unsupported adapter), the CLI warns and opens an empty atlas. + +An absent, unreadable, non-JSON, or wrong-format file is a usage error (exit 2). The +CLI checks the format tag and metadata object; the browser validates the pattern +array separately. The atlas currently requires at least one pattern, so an empty or +malformed array produces a browser error rather than a CLI exit-code failure. + +It shows the summary, filters by name, kernel and kind, groups by source file or by +kind, and expands each pattern's MLIR on demand. Nothing is uploaded; the file is read +in the browser. The page ships with the package, at +`winml/modelkit/commands/assets/cgc-pattern-atlas.html`, so `--open` finds it in an +installed wheel as well as in a checkout; open that file directly if you would rather +not go through the command. + +## Patterns and rules are not the same number + +A **pattern** is one `cgc_pattern.pattern` declaration. A **rule** is one flat match +alternative after `any_of` expansion: `any_of { all_of {…} all_of {…} }` contributes one +rule per branch, and several such blocks in one pattern multiply. + +They can diverge. For example, the following recorded counts show why both are useful: + +| driver | patterns | rules | +|---|---:|---:| +| 32.0.16.2009 | 13 | 13 | +| 32.0.16.2035 | 60 | 60 | +| 32.0.16.3004 | 51 | **63** | + +Reading only the declaration count across the 2035 → 3004 examples shows 60 → 51, +while the rule count increases from 60 to 63 because of `any_of` alternatives. +Neither number alone proves broader model coverage or better performance: these are +counts of declarations, not a compilation or execution test. + +Beware `apply_native_constraint "cgc_is_any_of"`, which is an op-family whitelist and +not a rule multiplier at all. A substring search for `any_of` conflates the two. + +## Finding the redist + +For a run that reaches redist resolution, an **explicitly named** redist -- the flag, +or the first non-empty environment variable -- is the only candidate. If it cannot +be used, the run exits 1 rather than answering about a different runtime. Only the +implicit `bin/` roots are tried in sequence, stopping at the first candidate that +passes the on-disk checks: + +| # | Source | Where | +|---|---|---| +| 1 | `--d3d12-dir PATH` | the flag; when given it is the **only** candidate | +| 2 | `$WINML_D3D12_DIR` | environment | +| 3 | `$D3D12_DIR` | environment | +| 4 | `/bin/` | the directory holding the virtual environment — beside `.venv` | +| 5 | `/bin/` | the directory holding `pyproject.toml`, for a source checkout | +| 6 | `/bin/` | the installed `winml` package | +| 7 | `/bin/` | the directory holding that package | +| 8 | `/bin/` | the environment root itself, e.g. `/bin/` | + +An explicit `--d3d12-dir` short-circuits the rest, so a wrong path fails loudly instead +of silently falling through to some other core. + +In PowerShell, set the environment variable with `$env:`: + +```powershell +$env:WINML_D3D12_DIR = 'C:\path\to\D3D12' +winml cgc adapters +``` + +### The `bin/` default + +A redist dropped in a `bin/` folder at the root of either a development worktree or a +wheel installation is used with no flag and no environment variable, from any working +directory. + +**Development environment** — the `bin/` folder sits **beside `.venv`**: + +``` +my-project/ + .venv/ + bin/D3D12/D3D12Core.dll <- dropped here +``` + +This is the rule that matters in practice, and it is checked first. It is derived from +`sys.prefix`, not from the source tree, so it keeps working once winml-cli is installed +from a wheel into a project's venv — at that point there is no `pyproject.toml` to walk +up to from `site-packages`, but `.venv`'s parent is still the project. + +A source checkout is also located by its `pyproject.toml`, which covers a venv kept +somewhere other than the worktree. In the usual case the two are the same directory and +collapse to one set of candidates. + +``` +winml-cli/ + .venv/ + bin/D3D12/D3D12Core.dll <- dropped here + pyproject.toml + src/winml/... +``` + +**Wheel installation** — these additional locations are searched for a locally +supplied redist: + +``` +/ + bin/D3D12/D3D12Core.dll <- + Lib/site-packages/ + bin/D3D12/D3D12Core.dll <- + winml/ + bin/D3D12/D3D12Core.dll <- +``` + +These are discovery locations, not a promise that the wheel contains an Agility SDK. +winml-cli does not download a redist; supply one separately. + +Each root is checked twice — `bin/D3D12/` first, then `bin/` holding `D3D12Core.dll` +directly. The `D3D12` subfolder is preferred because it is the Agility SDK's own +convention, mirroring the `D3D12SDKPath` an application exports, but a flat `bin/` works. + +Duplicate paths are checked once. Automatic worktree and virtual-environment-parent +roots at a filesystem root are excluded. + +Without a separately supplied redist, the `MLIR` column reads `?`. Which candidate was +chosen is on the `redist:` line whether or not `-v` was passed. `-v` adds each +adapter's raw capability HRESULT on stderr, and the negotiated IR version when +nonzero: + +```text + NVIDIA GeForce RTX 5060 Laptop GPU: feature 70 -> 0x00000000 (S_OK) + AMD Radeon(TM) 880M Graphics: feature 70 -> 0x80004001 (E_NOTIMPL) + Microsoft Basic Render Driver: feature 70 -> 0x8000FFFF (E_UNEXPECTED) +``` + +A candidate directory is accepted only after its `D3D12Core.dll` is found, its PE +machine type matches this interpreter, and its exported `D3D12SDKVersion` is read — by +parsing the export table, never by loading the DLL, which would pin a runtime before the +choice has been made. `isdir()` is not enough on its own. + +Passing these checks does not guarantee D3D12 can load the redist. If runtime loading +then fails, the command does not retry another candidate: an `adapters` listing with +an automatically discovered redist warns and shows `?`, while a named redist or a +`patterns` query fails with exit 1. + +### The SDK 720 floor + +A core older than SDK 720 is **never used**. It is *named* in the failure report when +no usable redist is found at all; when a later candidate does work, the rejection is not +reported -- the `redist:` line names the core that answered, and that is the one the +result describes: + +```console +$ winml cgc adapters --d3d12-dir C:\Windows\System32 +cgc: no usable D3D12 Agility SDK redist found (need a directory containing + D3D12Core.dll of SDK 720 or newer, built for x64). + +Looked in, in order: + 1. --d3d12-dir C:\Windows\System32 + SDK 616 is older than 720 and does not serve D3D MLIR programs +... +Error: --d3d12-dir named a redist that cannot be used; refusing to answer about a different one. +``` + +This exits **1**. + +This matters because an old core does not fail loudly — it answers `E_INVALIDARG` for +the feature query, which would report a perfectly capable driver as `MLIR no`. Falling +back to one would create a healthy-looking device with the MLIR path silently absent. +The CLI does not fall back to the inbox core. Candidate rejection is based on SDK +version and architecture, not on the directory name alone. + +The SDK version also selects the **wire ABI** used by this CLI: SDK 720 selects the +exchange by GUID; SDK 721 and newer use a `Type` enum plus a negotiated IR version. +There is no `--abi` flag: point `--d3d12-dir` at the redist you want and the ABI follows. +Use separate invocations to compare redists. + +## Exit codes + +| code | meaning | +|---|---| +| 0 | the CLI completed: includes an unsupported adapter, an empty declaration, no adapters found, `adapters` with support unknown, help, or a viewer launch request | +| 1 | a live query could not complete, a dump was refused or could not be written, or the installed atlas asset is missing; an `adapters` listing through an unusable auto-discovered redist instead warns, shows `?`, and exits 0 | +| 2 | bad arguments: an unknown flag, `--dump` without `-a`, an unmatched selector after redist resolution, or a file rejected by `--open` validation | + +## When something looks wrong + +**`MLIR` shows `?` for every adapter.** Either no redist was found, or D3D12 refused the +one that was (the reason is then on stderr, e.g. a preview redist with Developer Mode +off). Either way the question could not be asked, and the listing itself is still +correct. When no redist was found, `-v` shows every place cgc looked. + +**`MLIR` shows `no` on a GPU you expect to support it.** Check `-v` for a device-creation +failure or the capability-query result. `S_OK` alone does not mean support: the +returned support value must also be nonzero. For SDK 721 or newer, a zero negotiated +IR version is reported as `no` even if the HRESULT is `S_OK`. + +**The exchange fails after the driver claimed support.** The declaration was not +retrieved, so the run exits 1 rather than reporting the adapter as unsupported. +One possible cause is a driver/redist ABI mismatch: a driver can advertise capability +through a 721-shaped query but reject the exchange with `DXGI_ERROR_UNSUPPORTED`. +Check the HRESULT and try a redist compatible with that driver. + +**A driver answers in bytecode.** The CLI prints its byte count rather than pattern +counts. With `--dump`, it writes `patterns.mlirbc` and `metadata.txt`, but no +`patterns.json`. Rendering bytecode as text needs an MLIR-aware tool such as `cgc-opt`. + +**`--open` reports a missing atlas asset.** The installed package must contain the HTML +page. Reinstall a package that includes it. If `--dump` already completed, the files +remain available even though opening the viewer failed. diff --git a/docs/commands/overview.md b/docs/commands/overview.md index 0e9374e6e..4cd552628 100644 --- a/docs/commands/overview.md +++ b/docs/commands/overview.md @@ -6,7 +6,7 @@ shares a consistent invocation style — `winml [flags]` — and the same global flags are available on the root `winml` group. The commands group by user intent. **Discover** (`sys`, `inspect`, `catalog`, -`analyze`) helps you understand your hardware and model before writing any +`cgc`, `analyze`) helps you understand your hardware and model before writing any artifacts. **Configure** (`config`, `optimize`) produces a reusable build configuration and tunes the ONNX graph. **Build** (`export`, `quantize`, `compile`, `build`) runs the pipeline stages that produce deployment artifacts. @@ -26,6 +26,7 @@ measure speed and accuracy. | [`sys`](sys.md) | Discover | Inspect your machine — devices, EPs, and runtime versions at a glance. | | [`inspect`](inspect.md) | Discover | Inspect a model's tasks, classes, and hierarchy before committing to an export. | | [`catalog`](catalog.md) | Discover | Browse the curated winml-cli catalog of validated models and benchmarks. | +| [`cgc`](cgc.md) | Discover | Inspect D3D12 adapters and the MLIR patterns their drivers declare. | | [`config`](config.md) | Configure | Generate a reusable build configuration for a Hugging Face model or ONNX file. | | [`export`](export.md) | Build | Convert a PyTorch / Hugging Face model to ONNX, preserving module hierarchy. | | [`analyze`](analyze.md) | Build | Verify an ONNX model is compatible with a target execution provider before deployment. | @@ -41,6 +42,8 @@ measure speed and accuracy. - **I want to see what hardware and EPs I have** → `winml sys` - **I want to know if my model is supported** → `winml inspect` - **I want to browse validated models with known benchmarks** → `winml catalog` +- **I want to know whether a GPU driver implements D3D MLIR programs** → `winml cgc adapters` +- **I want to see the MLIR patterns a GPU driver declares** → `winml cgc patterns` - **I want to verify EP operator compatibility before compiling** → `winml analyze` - **I want to convert a Hugging Face model to ONNX** → `winml export` - **I want to run the whole pipeline in one go** → `winml build` diff --git a/mkdocs.yml b/mkdocs.yml index 262d1c577..63d33026c 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -114,6 +114,7 @@ nav: - sys: commands/sys.md - inspect: commands/inspect.md - catalog: commands/catalog.md + - cgc: commands/cgc.md - Configure: - config: commands/config.md - Build: diff --git a/pyproject.toml b/pyproject.toml index 240dd6c7e..112d41589 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -176,6 +176,9 @@ where = [ "src" ] include = [ "winml", "winml.*" ] [tool.setuptools.package-data] +# The CGC pattern atlas is opened by `winml cgc patterns --open`, so it ships in +# the wheel. It is not a Python file, so packages.find never picks it up. +"winml.modelkit.commands" = [ "assets/*.html" ] "winml.modelkit" = [ "py.typed" ] "winml.modelkit.data" = [ "hub_models.json" ] "winml.modelkit.analyze" = [ diff --git a/src/winml/modelkit/commands/assets/cgc-pattern-atlas.html b/src/winml/modelkit/commands/assets/cgc-pattern-atlas.html new file mode 100644 index 000000000..776087364 --- /dev/null +++ b/src/winml/modelkit/commands/assets/cgc-pattern-atlas.html @@ -0,0 +1,608 @@ + + + + + + +CGC Pattern Atlas + + + +
+
+

CGC · driver-declared patterns

+

CGC Pattern Atlas

+
+
+ +
+ Drop a patterns.json anywhere on the page, or + + + +

+
+
+
+ +
+
+ +
+
+ group + + +
+ +
+
+ +
+

No patterns loaded yet. Drop a patterns.json written by + winml cgc patterns -a <adapter> --dump anywhere on this page.

+
+
+ +
+ +
+
+ Generated from a patterns.json written by winml cgc patterns --dump. + Kind is read from each pattern's body, never its name. +
+
+ + + + + + + diff --git a/src/winml/modelkit/commands/cgc.py b/src/winml/modelkit/commands/cgc.py new file mode 100644 index 000000000..01684cfd7 --- /dev/null +++ b/src/winml/modelkit/commands/cgc.py @@ -0,0 +1,1986 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""D3D12 adapters on this machine, and the patterns their drivers declare. + +Two questions get asked on every new machine and after every driver update: which +adapters are here, and what does this driver's MLIR-program implementation claim it +can match. Both answers come from public Windows APIs through ``ctypes``, so nothing +has to be built and no extra package has to be installed. + +It does need one thing: an Agility SDK ``D3D12Core.dll`` of SDK version 720 or newer. +The MLIR-program feature is a preview feature and the inbox D3D12 runtime does not +serve it. Which redist answered opens every run on the ``redist:`` line, because the +same driver reports different things through the 720 and 721 exchange shapes; if none +can be found the failure names every place that was searched. + +Columns: + IDX selection index, as accepted by ``-a`` + ADAPTER DXCore DriverDescription + DRIVER DXCore DriverVersion, the four 16-bit parts + TYPE hardware, integrated or software + ATTRIBUTES the DXCore attributes the adapter advertises: ML, CC, GFX + MLIR whether the driver implements the MLIR-program exchange; ``?`` when no + Agility SDK redist was found, since the question cannot be asked + +Counts are reported as ``N patterns / M rules``. A pattern is one +``cgc_pattern.pattern`` declaration; a rule is one flat match alternative after +``any_of`` expansion. The two differ, and comparing declaration counts alone across +drivers can read a loss where the matching surface actually grew. +""" + +from __future__ import annotations + +import ctypes +import ctypes.wintypes +import json +import logging +import os +import re +import struct +import sys +import sysconfig +import tempfile +import uuid +import webbrowser +from dataclasses import dataclass +from importlib import resources +from pathlib import Path +from typing import TYPE_CHECKING, Any, cast + +import click + +from .. import commands +from ..utils import cli as cli_utils +from ..utils.logging import configure_logging + + +if TYPE_CHECKING: + from collections.abc import Callable, Iterator + +logger = logging.getLogger(__name__) + + +class GUID(ctypes.Structure): + """A Windows GUID, laid out for ``ctypes``.""" + + _fields_ = ( + ("Data1", ctypes.c_uint32), + ("Data2", ctypes.c_uint16), + ("Data3", ctypes.c_uint16), + ("Data4", ctypes.c_ubyte * 8), + ) + + +def guid(text: str) -> GUID: + """Build a :class:`GUID` from its registry form, as the Windows SDK headers print it. + + Args: + text: For example ``78EE5945-C36E-4B13-A669-005DD11C0F06``. + + Returns: + The populated GUID. + """ + return GUID.from_buffer_copy(uuid.UUID(text).bytes_le) + + +def vcall(this: ctypes.c_void_p, index: int, restype: Any, *argtypes: Any) -> Callable[..., Any]: + """Bind vtable slot *index* of COM pointer *this*. + + Replaces a comtypes dependency with six lines. + + Args: + this: The COM interface pointer. + index: Zero-based vtable slot. + restype: ``ctypes`` return type. + *argtypes: ``ctypes`` argument types, excluding the implicit ``this``. + + Returns: + A callable that invokes the method, passing ``this`` automatically. + """ + vtbl = ctypes.cast(this, ctypes.POINTER(ctypes.POINTER(ctypes.c_void_p)))[0] + proto = ctypes.WINFUNCTYPE(restype, ctypes.c_void_p, *argtypes) + fn = proto(vtbl[index]) + + return lambda *args: fn(this, *args) + + +def release(p: ctypes.c_void_p | None) -> None: + """Call ``IUnknown::Release`` on *p* when it is non-NULL. + + Args: + p: The COM interface pointer, or None. + """ + if p: + vcall(p, 2, ctypes.c_ulong)() + + +E_FAIL = -2147467259 # 0x80004005 as a signed HRESULT + +HRESULTS = { + 0x00000000: "S_OK", + 0x80004001: "E_NOTIMPL", + 0x80004005: "E_FAIL", + 0x8000FFFF: "E_UNEXPECTED", + 0x80070057: "E_INVALIDARG", + 0x887A0004: "DXGI_ERROR_UNSUPPORTED", + 0x887E0001: "D3D12_ERROR_ADAPTER_NOT_FOUND", + 0x887E0002: "D3D12_ERROR_DRIVER_VERSION_MISMATCH", + 0x887E0003: "D3D12_ERROR_INVALID_REDIST", +} + + +def hrs(code: int) -> str: + """Render an HRESULT as hex, plus its name when one is known. + + Args: + code: The HRESULT, signed or unsigned. + + Returns: + For example ``0x80070057 (E_INVALIDARG)``; unknown codes print raw. + """ + u = code & 0xFFFFFFFF + return f"0x{u:08X}" + (f" ({HRESULTS[u]})" if u in HRESULTS else "") + + +def hr(name: str, code: int) -> None: + """Exit when *code* is a failing HRESULT. + + Args: + name: The call being checked, for the message. + code: The returned HRESULT. + + Raises: + click.ClickException: When *code* indicates failure. + """ + if code < 0: + raise click.ClickException(f"{name} failed: {hrs(code)}") + + +# ----------------------------------------------------------------------- dxcore.dll +# dxcore_interface.h. IDXCoreAdapterFactory : IUnknown -> 3 CreateAdapterList. +# IDXCoreAdapterList -> 3 GetAdapter, 4 GetAdapterCount, 7 Sort. +# IDXCoreAdapter -> 4 IsAttributeSupported, 5 IsPropertySupported, 6 GetProperty, +# 7 GetPropertySize. +IID_IDXCoreAdapterFactory = guid("78EE5945-C36E-4B13-A669-005DD11C0F06") +IID_IDXCoreAdapterList = guid("526C7776-40E9-459B-B711-F32AD76DFC28") +IID_IDXCoreAdapter = guid("F0DB4C7F-FE5A-42A2-BD62-F2A6CF6FC83E") +ATTR_GENERIC_ML = guid("B71B0D41-1088-422F-A27C-0250B7D3A988") +ATTR_CORE_COMPUTE = guid("248E2800-A793-4724-ABAA-23A6DE1BE090") +ATTR_D3D12_GRAPHICS = guid("0C9ECE4D-2F6E-4F01-8C96-E89E331B47B1") + +PROP_DRIVER_VERSION, PROP_DRIVER_DESCRIPTION = 1, 2 +PROP_IS_HARDWARE, PROP_IS_INTEGRATED = 11, 12 +PREF_HARDWARE, PREF_HIGH_PERFORMANCE = 0, 2 + +Adapter = dict[str, Any] + + +def read_adapter_props(ad: ctypes.c_void_p) -> Adapter: + """Read every property the listing shows from one ``IDXCoreAdapter``. + + ``IsPropertySupported`` returns a C++ bool, one byte. + + Args: + ad: The adapter interface pointer. + + Returns: + The adapter record, without ``index``, ``ptr``, ``mlir`` or ``ir_version``. + """ + is_prop = vcall(ad, 5, ctypes.c_bool, ctypes.c_uint32) + get_prop = vcall(ad, 6, ctypes.c_long, ctypes.c_uint32, ctypes.c_size_t, ctypes.c_void_p) + get_size = vcall(ad, 7, ctypes.c_long, ctypes.c_uint32, ctypes.POINTER(ctypes.c_size_t)) + is_attr = vcall(ad, 4, ctypes.c_bool, ctypes.POINTER(GUID)) + + def read(pid: int, label: str) -> Any: + # The size DXCore reports is honoured rather than assumed. The two flags are 1 + # byte, not a 4-byte BOOL; reading them as BOOL marks every adapter "software", + # which silently changes which adapter the default -a selects. + n = ctypes.c_size_t() + hr(f"GetPropertySize({label})", get_size(pid, ctypes.byref(n))) + buf = ctypes.create_string_buffer(n.value) + hr(f"GetProperty({label})", get_prop(pid, n.value, buf)) + return buf + + p: Adapter = { + "description": "", + "driver_version": "", + "is_hardware": False, + "is_integrated": False, + } + if is_prop(PROP_DRIVER_DESCRIPTION): + # The size includes the NUL, which .value stops at. + description = read(PROP_DRIVER_DESCRIPTION, "DriverDescription").value + p["description"] = description.decode("utf-8", "replace") + if is_prop(PROP_DRIVER_VERSION): + v = ctypes.c_uint64() + hr("GetProperty(DriverVersion)", get_prop(PROP_DRIVER_VERSION, 8, ctypes.byref(v))) + p["driver_version"] = format_driver_version(v.value) + for key, pid in (("is_hardware", PROP_IS_HARDWARE), ("is_integrated", PROP_IS_INTEGRATED)): + if is_prop(pid): + p[key] = any(read(pid, key).raw) + p["generic_ml"] = bool(is_attr(ctypes.byref(ATTR_GENERIC_ML))) + p["core_compute"] = bool(is_attr(ctypes.byref(ATTR_CORE_COMPUTE))) + p["d3d12_graphics"] = bool(is_attr(ctypes.byref(ATTR_D3D12_GRAPHICS))) + return p + + +def format_driver_version(raw: int) -> str: + """Render the DXCore u64 driver version as four 16-bit parts, high word first. + + Args: + raw: The packed 64-bit value. + + Returns: + For example ``32.0.16.3004``. + """ + return f"{(raw >> 48) & 0xFFFF}.{(raw >> 32) & 0xFFFF}.{(raw >> 16) & 0xFFFF}.{raw & 0xFFFF}" + + +def list_adapters() -> list[Adapter]: + """Enumerate every D3D12 adapter, sorted hardware-then-high-performance. + + Creates no device, so this costs nothing and needs no redist. + + Returns: + One record per adapter. Each owns a ``ptr`` that :func:`close_adapters` + must release. + """ + try: + dxcore = ctypes.WinDLL("dxcore.dll") + except OSError as e: + raise click.ClickException(f"dxcore.dll could not be loaded: {e}") from e + dxcore.DXCoreCreateAdapterFactory.restype = ctypes.c_long + dxcore.DXCoreCreateAdapterFactory.argtypes = [ + ctypes.POINTER(GUID), + ctypes.POINTER(ctypes.c_void_p), + ] + + factory = ctypes.c_void_p() + hr( + "DXCoreCreateAdapterFactory", + dxcore.DXCoreCreateAdapterFactory( + ctypes.byref(IID_IDXCoreAdapterFactory), ctypes.byref(factory) + ), + ) + create_list = vcall( + factory, + 3, + ctypes.c_long, + ctypes.c_uint32, + ctypes.POINTER(GUID), + ctypes.POINTER(GUID), + ctypes.POINTER(ctypes.c_void_p), + ) + + alist = ctypes.c_void_p() + hr( + "CreateAdapterList(GENERIC_ML)", + create_list( + 1, + ctypes.byref(ATTR_GENERIC_ML), + ctypes.byref(IID_IDXCoreAdapterList), + ctypes.byref(alist), + ), + ) + if vcall(alist, 4, ctypes.c_uint32)() == 0: + release(alist) + alist = ctypes.c_void_p() + hr( + "CreateAdapterList(CORE_COMPUTE)", + create_list( + 1, + ctypes.byref(ATTR_CORE_COMPUTE), + ctypes.byref(IID_IDXCoreAdapterList), + ctypes.byref(alist), + ), + ) + + prefs = (ctypes.c_uint32 * 2)(PREF_HARDWARE, PREF_HIGH_PERFORMANCE) + vcall(alist, 7, ctypes.c_long, ctypes.c_uint32, ctypes.POINTER(ctypes.c_uint32))(2, prefs) + + count = vcall(alist, 4, ctypes.c_uint32)() + get_adapter = vcall( + alist, + 3, + ctypes.c_long, + ctypes.c_uint32, + ctypes.POINTER(GUID), + ctypes.POINTER(ctypes.c_void_p), + ) + adapters: list[Adapter] = [] + for i in range(count): + ad = ctypes.c_void_p() + hr("GetAdapter", get_adapter(i, ctypes.byref(IID_IDXCoreAdapter), ctypes.byref(ad))) + props = read_adapter_props(ad) + props["index"] = i + props["ptr"] = ad + props["mlir"], props["ir_version"] = None, 0 # set by probe_mlir_support + adapters.append(props) + release(alist) + release(factory) + return adapters + + +def close_adapters(adapters: list[Adapter]) -> None: + """Release every adapter pointer in *adapters*. + + Args: + adapters: The records returned by :func:`list_adapters`. + """ + for a in adapters: + release(a.pop("ptr", None)) + + +def adapter_slug(description: str) -> str: + """Turn an adapter description into a directory-safe name. + + ``"NVIDIA GeForce RTX 5090 D"`` becomes ``"nvidia-geforce-rtx-5090-d"``. + + Args: + description: The DXCore DriverDescription. + + Returns: + The slug. + """ + s = re.sub(r"\((?:r|tm|c)\)", "", description, flags=re.I) + s = re.sub(r"[^0-9a-z]+", "-", s.lower()) + # Never empty: dump_dir joins it, and "" would collapse the per-adapter level. + return s.strip("-") or "adapter" + + +def select_adapter(adapters: list[Adapter], wanted: str) -> Adapter: + """Resolve ``-a`` to one adapter. + + ``-a`` takes a description substring or a bare index. With neither, the first + hardware adapter wins -- not index 0, which on many machines is a software adapter. + + Args: + adapters: Every adapter found; never empty. + wanted: The ``-a`` value, possibly empty. + + Returns: + The selected adapter. + + Raises: + click.UsageError: When nothing matches *wanted*. + """ + if not wanted: + return next((a for a in adapters if a["is_hardware"]), adapters[0]) + # A number is an index when one exists, else a substring: "5060" names the RTX + # 5060 on a three-adapter machine rather than failing as an out-of-range index. + if wanted.isdigit() and int(wanted) < len(adapters): + return adapters[int(wanted)] + hit = next((a for a in adapters if wanted.lower() in a["description"].lower()), None) + if hit is None: + found = ", ".join(a["description"] for a in adapters) + raise click.UsageError(f"no adapter matching {wanted!r}; found: {found}") + return hit + + +# ------------------------------------------------------- the Agility SDK redist +# The MLIR-program feature is a preview feature: the inbox D3D12 runtime answers +# E_INVALIDARG for it, so a device must be created against an Agility SDK redist. Which +# exchange ABI to speak is a property of that redist, which is why its SDK version is +# read here rather than guessed, and read by parsing the PE export table rather than by +# loading the DLL -- loading pins a core in-process before the choice has been made. + +MACHINES = {0x014C: "x86", 0x8664: "x64", 0xAA64: "arm64"} + +#: The oldest Agility SDK that serves D3D MLIR programs. Anything below this answers +#: E_INVALIDARG for the feature and would report a capable driver as "no" -- a silently +#: wrong answer, so such a core is named and skipped rather than used. +MIN_SDK_VERSION = 720 + +#: Environment variables naming a redist directory, in the order they are consulted. +REDIST_ENV_VARS = ("WINML_D3D12_DIR", "D3D12_DIR") + +#: Where a redist may sit under each search root. The Agility SDK's own convention is a +#: ``D3D12`` subfolder beside the binaries; ``bin`` holding ``D3D12Core.dll`` directly +#: is accepted too. +REDIST_SUBDIRS = ("bin/D3D12", "bin") + + +def host_machine() -> str: + """Return this interpreter's machine name as the PE header spells it. + + Returns: + ``x64``, ``x86``, ``arm64``, or ``?``. + """ + # The interpreter's own build, which is what decides whether this process can load + # the DLL. platform.machine() would answer the same on a native install, but on 3.11 + # it runs `cmd /c ver` first, ~45 ms on every run. + return {"win-amd64": "x64", "win32": "x86", "win-arm64": "arm64"}.get( + sysconfig.get_platform(), "?" + ) + + +def pe_export_u32(path: Path, want: bytes) -> tuple[str | None, int | None]: + """Read a UINT32 data export out of a PE file without loading it. + + Args: + path: The PE file. + want: The export name, as bytes. + + Returns: + ``(machine, value)``, either of which is None when unavailable. + + Raises: + struct.error: When the file is truncated, or the export table points outside + every section. + """ + data = path.read_bytes() + if data[:2] != b"MZ": + return None, None + (pe,) = struct.unpack_from(" int: + for vsize, vaddr, rawsize, rawptr in sections: + if vaddr <= rva < vaddr + max(vsize, rawsize): + return int(rawptr + (rva - vaddr)) + # A malformed export table is a reason to reject this candidate and try + # the next one, so raise something inspect_redist already catches. + raise struct.error(f"RVA 0x{rva:X} is outside every section") + + n_names, a_funcs, a_names, a_ords = struct.unpack_from(" tuple[int | None, str]: + """Decide whether *path* is a usable Agility SDK redist. + + Args: + path: A candidate directory. + + Returns: + ``(sdk_version, reason)``. ``sdk_version`` is None when the directory cannot + be used, and ``reason`` then says why. + """ + dll = path / "D3D12Core.dll" + try: + # Every filesystem call is inside the guard: an unreachable share, an + # ACL-restricted directory or a stale mount must be a rejected candidate, not a + # traceback, because every implicit root is probed on every run. + if not path.is_dir(): + return None, "missing" + if not dll.is_file(): + return None, "empty" if not any(path.iterdir()) else "no D3D12Core.dll" + arch, version = pe_export_u32(dll, b"D3D12SDKVersion") + except (OSError, struct.error) as e: + return None, f"unreadable: {e}" + if version is None: + return None, "no D3D12SDKVersion export" + if arch != host_machine(): + return None, f"arch {arch} != {host_machine()}" + if version < MIN_SDK_VERSION: + return None, ( + f"SDK {version} is older than {MIN_SDK_VERSION} and does not serve D3D MLIR programs" + ) + return version, f"SDK {version}" + + +def _redist_roots() -> list[tuple[str, Path]]: + """Return the directories that may hold a ``bin`` folder, in search order. + + Two placements are supported, matching where a redist is actually dropped: + the root of a development environment, and the root of a wheel installation. + + The development root is the directory holding the virtual environment -- the + ``bin`` folder sits beside ``.venv``. That is checked before the worktree + marker because it is the only rule that works once winml-cli is installed + from a wheel into a project's venv: there is no ``pyproject.toml`` to walk up + to from ``site-packages``, but ``.venv``'s parent is still the project. + + "Installation root" is deliberately read broadly -- the installed ``winml`` + package, the ``site-packages`` directory holding it, and the environment + prefix -- because a wheel can deliver files to any of the three depending on + whether it ships them as package data or as data files. + + Returns: + ``(label, root)`` pairs. Duplicates are dropped downstream, so the common + case where the venv sits in the worktree yields one set of candidates. + """ + package = Path(__file__).resolve().parent.parent.parent # .../winml + roots: list[tuple[str, Path]] = [] + # sys.prefix differs from base_prefix only inside a virtual environment; outside + # one its parent is the interpreter's install root and means nothing here. + # + # This is the one root not bounded by the project or install tree -- it walks *up* + # out of the venv -- so it stops at a filesystem anchor. A venv placed directly at + # a root would otherwise make that root the highest-priority candidate: `C:\.venv` + # yields `C:\bin\D3D12` and `\\server\share\venv` yields `\\server\share\bin`. + # Since the resolved directory is handed to CreateDeviceFactory, which loads + # D3D12Core.dll as native code into this process, and default Windows ACLs let a + # standard user create directories at a drive root, such a layout would turn the + # first candidate into a plantable DLL path. resolve() first so a junctioned + # prefix does not defeat the de-duplication below either. + if sys.prefix != sys.base_prefix: + venv_parent = Path(sys.prefix).resolve().parent + if venv_parent != venv_parent.parent: + roots.append(("", venv_parent)) + # A development checkout: the nearest ancestor holding a pyproject.toml; none in a + # wheel install. Same anchor guard as : the walk reaches every + # ancestor, so a pyproject.toml at a volume root would otherwise become a base. + here = Path(__file__).resolve() + repo = next((d for d in here.parents if (d / "pyproject.toml").is_file()), None) + if repo is not None and repo != repo.parent: + roots.append(("", repo)) + roots.append(("", package)) + roots.append(("", package.parent)) + roots.append(("", Path(sys.prefix))) + return roots + + +def _named_redist(explicit: str = "") -> tuple[str, Path] | None: + """Return the redist the caller named by hand, if any. + + The single place that decides precedence -- the flag, then each variable in + :data:`REDIST_ENV_VARS`, first non-empty wins -- so the candidate list and the + decision whether an unusable redist is fatal can never disagree. + + Args: + explicit: The ``--d3d12-dir`` value, possibly empty. + + Returns: + ``(label, path)``, or None when the caller named nothing. + """ + if explicit: + return "--d3d12-dir", Path(explicit) + for var in REDIST_ENV_VARS: + value = os.environ.get(var) + if value: + return f"${var}", Path(value) + return None + + +# TODO(tests): the documented search order (flag, variables, then each bin root) +# and the filesystem-anchor guard on have no test. +def redist_candidates(explicit: str = "") -> Iterator[tuple[str, Path]]: + """Yield ``(label, path)`` in resolution order. + + A redist the caller named (see :func:`_named_redist`) is the only candidate, usable + or not: falling through to a different core would answer the capability question + about a runtime the caller did not ask for. + + Args: + explicit: The ``--d3d12-dir`` value, possibly empty. + + Yields: + A label for the failure report, and the directory to inspect. + """ + named = _named_redist(explicit) + if named is not None: + yield named + return + seen: set[Path] = set() + for label, root in _redist_roots(): + for sub in REDIST_SUBDIRS: + path = root / sub + if path not in seen: + seen.add(path) + yield f"{label}/{sub}", path + + +def resolve_redist( + explicit: str = "", +) -> tuple[Path | None, int | None, list[tuple[str, Path, str]]]: + """Find the first usable redist. + + Args: + explicit: The ``--d3d12-dir`` value, possibly empty. + + Returns: + ``(path, sdk_version, tried)``. *path* is None when nothing usable was found, + and *tried* records every candidate and why it was rejected. + """ + tried: list[tuple[str, Path, str]] = [] + for label, path in redist_candidates(explicit): + version, reason = inspect_redist(path) + tried.append((label, path, reason)) + if version is not None: + return path, version, tried + return None, None, tried + + +def redist_failure(tried: list[tuple[str, Path, str]]) -> str: + """Name every place looked and why each was rejected. + + Silence here costs a day. + + Args: + tried: The candidates from :func:`resolve_redist`. + + Returns: + The multi-line report, for stderr. + """ + out = [ + "cgc: no usable D3D12 Agility SDK redist found (need a directory containing", + f" D3D12Core.dll of SDK {MIN_SDK_VERSION} or newer, built for {host_machine()}).", + "", + "Looked in, in order:", + ] + for i, (label, path, reason) in enumerate(tried, 1): + out.append(f" {i}. {label:<26} {path}") + out.append(f" {reason}") + out += [ + "", + "Fix one of:", + " - point at one: winml cgc adapters --d3d12-dir ", + f" - set the env: {REDIST_ENV_VARS[0]}=", + "", + "Note: C:/Windows/System32/D3D12Core.dll is the inbox core and does not", + "implement D3D MLIR programs. cgc does not fall back to it, because doing so", + "creates a healthy-looking device with the MLIR path silently absent.", + ] + return "\n".join(out) + + +# ------------------------------------------------------------------- d3d12.dll +CLSID_D3D12SDKConfiguration = guid("7CDA6ACA-A03E-49C8-9458-0334D20E07CE") +IID_ID3D12SDKConfiguration1 = guid("8AAF9303-AD25-48B9-9A57-D9C37E009D9F") +IID_ID3D12DeviceFactory = guid("61F307D3-D34E-4E7C-8374-3BA4DE23CCCB") +IID_ID3D12Device = guid("189819F1-1DB6-4B57-BE54-1821339B85F7") + +FEATURE_LEVELS = (0x0100, 0x1000, 0xB000) # 1_0_GENERIC, 1_0_CORE, 11_0 + +FEATURE_MLIR_EXCHANGE = 69 # the same id under both ABIs +FEATURE_MLIR_70 = 70 # INTERFACE_SUPPORT on 720, COMPUTE_GRAPH_VERSION on 721 +MLIR_EXCHANGE_SUBGRAPH_DECLARATION = 0 + +# include/utilities/MlirInterfaceGuids.h:51 -- the one DXML-private constant here. +GUID_SUBGRAPH_DECLARATION_REQUEST = guid("EB53032A-1116-4E71-8309-54347C1E5A26") + +CGC_MAX_IR_VERSION = (0, 7, 0, 0) # cmake/version.cmake CGC_VERSION; an upper bound + + +class FeatureDataMLIRComputeGraphVersion(ctypes.Structure): + """``D3D12_FEATURE_DATA_MLIR_COMPUTE_GRAPH_VERSION``, the 721 capability query.""" + + _fields_ = (("HighestVersion", ctypes.c_uint64),) + + +class FeatureDataMLIRExchange721(ctypes.Structure): + """``D3D12_FEATURE_DATA_MLIR_EXCHANGE`` as SDK 721 shapes it. + + 48 bytes on x64, with 4 bytes of padding after Type. The driver reads it by size, + so a drifted layout is silently wrong; ``test_721_exchange_struct_is_48_bytes`` + pins it. + """ + + _fields_ = ( + ("Type", ctypes.c_uint32), + ("IRVersion", ctypes.c_uint64), # D3D12_VERSION_NUMBER, a UINT64 union + ("pInputData", ctypes.c_void_p), + ("InputDataSizeInBytes", ctypes.c_size_t), + ("pOutputData", ctypes.c_void_p), + ("OutputDataSizeInBytes", ctypes.POINTER(ctypes.c_size_t)), + ) + + +class FeatureDataMLIRExchange720(ctypes.Structure): + """``D3D12_FEATURE_DATA_MLIR_EXCHANGE`` as SDK 720 shapes it.""" + + _fields_ = ( + ("MlirInterface", GUID), + ("pInputData", ctypes.c_void_p), + ("InputDataSizeInBytes", ctypes.c_size_t), + ("pOutputData", ctypes.c_void_p), + ("OutputDataSizeInBytes", ctypes.POINTER(ctypes.c_size_t)), + ) + + +class FeatureDataMLIRInterfaceSupport720(ctypes.Structure): + """``D3D12_FEATURE_DATA_MLIR_INTERFACE_SUPPORT``, the 720 capability query.""" + + _fields_ = ( + ("NumMlirInterfaces", ctypes.c_uint32), + ("pMlirInterfacesRequested", ctypes.POINTER(GUID)), + ("pMlirInterfacesSupported", ctypes.POINTER(ctypes.wintypes.BOOL)), + ) + + +def make_version_number(major: int, minor: int, build: int, rev: int) -> int: + """Pack four 16-bit parts into a ``D3D12_VERSION_NUMBER``. + + Args: + major: High word. + minor: Second word. + build: Third word. + rev: Low word. + + Returns: + The packed 64-bit value. + """ + return ( + ((major & 0xFFFF) << 48) + | ((minor & 0xFFFF) << 32) + | ((build & 0xFFFF) << 16) + | (rev & 0xFFFF) + ) + + +class RedistUnusable(click.ClickException): + """D3D12 cannot load, or rejects, the redist itself, so no adapter can be asked.""" + + +def create_device( + adapter_ptr: ctypes.c_void_p, redist: Path, sdk_version: int +) -> tuple[ctypes.c_void_p | None, int]: + """Create an ``ID3D12Device`` on *adapter_ptr* through the given redist. + + Args: + adapter_ptr: The DXCore adapter pointer. + redist: The directory holding ``D3D12Core.dll``. + sdk_version: That core's ``D3D12SDKVersion``. + + Returns: + ``(device, 0)``, or ``(None, hresult)`` when this adapter could not get a + device at any feature level. + + Raises: + RedistUnusable: When D3D12 cannot load or rejects the redist itself, since + then no adapter can be asked through it. + """ + # The path must be absolute: a relative one resolves against the host exe, which is + # python.exe. CreateDeviceFactory takes it as a narrow string, but D3D12 decodes it + # as UTF-8 -- not the active ANSI codepage. Encoding it as ANSI makes any path with + # a non-ASCII character (an accented profile name, CJK) fail with + # D3D12_ERROR_INVALID_REDIST. + # TODO(tests): no test covers a non-ASCII or CJK redist path. + try: + redist_utf8 = str(redist.resolve()).encode("utf-8") + except UnicodeEncodeError as e: # only an unpaired surrogate in a file name + raise click.ClickException(f"redist path is not valid Unicode: {redist}") from e + try: + d3d12 = ctypes.WinDLL("d3d12.dll") + except OSError as e: + # A host without the D3D12 runtime (Server Core without the graphics + # feature, some container images) must fail the way every other + # could-not-ask failure does, not with a bare traceback. + raise RedistUnusable(f"d3d12.dll could not be loaded: {e}") from e + d3d12.D3D12GetInterface.restype = ctypes.c_long + d3d12.D3D12GetInterface.argtypes = [ + ctypes.POINTER(GUID), + ctypes.POINTER(GUID), + ctypes.POINTER(ctypes.c_void_p), + ] + + cfg, factory = ctypes.c_void_p(), ctypes.c_void_p() + try: + rc = d3d12.D3D12GetInterface( + ctypes.byref(CLSID_D3D12SDKConfiguration), + ctypes.byref(IID_ID3D12SDKConfiguration1), + ctypes.byref(cfg), + ) + if rc < 0 or not cfg: + # This system cannot load an Agility core at all, so no adapter can be asked. + raise RedistUnusable( + f"D3D12 cannot load an Agility SDK redist on this system " + f"(D3D12GetInterface: {hrs(rc if rc < 0 else E_FAIL)})" + ) + + # ID3D12SDKConfiguration1 -> 4 CreateDeviceFactory. + create_factory = vcall( + cfg, + 4, + ctypes.c_long, + ctypes.c_uint32, + ctypes.c_char_p, + ctypes.POINTER(GUID), + ctypes.POINTER(ctypes.c_void_p), + ) + rc = create_factory( + sdk_version, + redist_utf8, + ctypes.byref(IID_ID3D12DeviceFactory), + ctypes.byref(factory), + ) + if rc < 0 or not factory: + # D3D12 rejected the redist itself, so no adapter can be asked through it. + # Returning per adapter would record every driver as "no" -- a confident + # wrong answer, and a permanent status=unsupported baseline under --dump. + raise RedistUnusable( + f"D3D12 rejected the redist at {redist} " + f"(CreateDeviceFactory: {hrs(rc if rc < 0 else E_FAIL)}); cannot ask " + f"whether any driver implements MLIR programs" + ) + + create = vcall( + factory, + 9, + ctypes.c_long, + ctypes.c_void_p, + ctypes.c_uint32, + ctypes.POINTER(GUID), + ctypes.POINTER(ctypes.c_void_p), + ) + last = 0 + for level in FEATURE_LEVELS: + dev = ctypes.c_void_p() + last = create(adapter_ptr, level, ctypes.byref(IID_ID3D12Device), ctypes.byref(dev)) + if last >= 0 and dev: + return dev, 0 + # No feature level gave this adapter a device: a per-adapter "no". A success + # code with a NULL device is still a failure, so it is never reported as S_OK. + return None, last if last < 0 else E_FAIL + finally: + release(factory) + release(cfg) + + +def check_feature_support(device: ctypes.c_void_p) -> Any: + """Bind ``ID3D12Device::CheckFeatureSupport``, vtable slot 13. + + Args: + device: The device pointer. + + Returns: + The bound callable. This is the only call that does real work. + """ + return vcall(device, 13, ctypes.c_long, ctypes.c_uint32, ctypes.c_void_p, ctypes.c_uint32) + + +def supports_exchange(device: ctypes.c_void_p, sdk_version: int) -> tuple[bool, int, int]: + """Ask whether this *driver* implements the exchange. + + Decided by value, never by HRESULT. + + Args: + device: The device pointer. + sdk_version: The redist's SDK version, which picks the ABI shape. + + Returns: + ``(supported, hresult, negotiated_ir_version)``. + """ + cfs = check_feature_support(device) + if sdk_version >= 721: + data = FeatureDataMLIRComputeGraphVersion(make_version_number(*CGC_MAX_IR_VERSION)) + rc = cfs(FEATURE_MLIR_70, ctypes.byref(data), ctypes.sizeof(data)) + # 70 is answered by the runtime, not the driver: WARP and Intel both return + # S_OK here with HighestVersion == 0. The value is the answer. + if rc < 0: + return False, rc, 0 + return data.HighestVersion != 0, rc, data.HighestVersion + supported = ctypes.wintypes.BOOL(0) + data_720 = FeatureDataMLIRInterfaceSupport720( + 1, ctypes.pointer(GUID_SUBGRAPH_DECLARATION_REQUEST), ctypes.pointer(supported) + ) + rc = cfs(FEATURE_MLIR_70, ctypes.byref(data_720), ctypes.sizeof(data_720)) + return (rc >= 0 and bool(supported.value)), rc, 0 + + +# TODO(tests): TestCapabilityOnRealHardware cannot fail -- it only asserts that +# mlir is a bool, which holds even when no device was created at all. +def probe_mlir_support( + adapter: Adapter, redist: Path, sdk_version: int, *, verbose: bool = False +) -> ctypes.c_void_p | None: + """Set ``adapter["mlir"]``, and ``adapter["ir_version"]`` when the driver answers. + + Any failing HRESULT means no exchange -- do not read into it. + + Args: + adapter: The adapter record, mutated in place. + redist: The directory holding ``D3D12Core.dll``. + sdk_version: That core's SDK version. + verbose: Write the raw capability HRESULT to stderr. + + Returns: + The device when the driver implements the exchange, for the caller to ask and + then release; else None. + + Raises: + RedistUnusable: From :func:`create_device`, when the redist itself cannot be + used; no adapter's answer is recorded then. + """ + device, rc = create_device(adapter["ptr"], redist, sdk_version) + if device is None: + adapter["mlir"] = False + if verbose: + click.echo(f" {adapter['description']}: CreateDevice failed: {hrs(rc)}", err=True) + return None + ok, rc, negotiated = supports_exchange(device, sdk_version) + adapter["mlir"], adapter["ir_version"] = ok, negotiated + if verbose: + extra = f", IR version {format_driver_version(negotiated)}" if negotiated else "" + click.echo(f" {adapter['description']}: feature 70 -> {hrs(rc)}{extra}", err=True) + if not ok: + release(device) + return None + return device + + +def mlir_exchange( + device: ctypes.c_void_p, sdk_version: int, ir_version: int +) -> tuple[bytes | None, int]: + """Fetch the subgraph-declaration payload. + + Uses the two-call size-then-fill protocol. + + Args: + device: The device pointer. + sdk_version: The redist's SDK version, which picks the ABI shape. + ir_version: The IR version feature 70 negotiated, for the 721 shape. + + Returns: + ``(payload, 0)``, or ``(None, hresult)`` when the exchange failed. + """ + cfs = check_feature_support(device) + + def build(out_ptr: ctypes.c_void_p | None, size_ptr: Any) -> ctypes.Structure: + if sdk_version >= 721: + return FeatureDataMLIRExchange721( + Type=MLIR_EXCHANGE_SUBGRAPH_DECLARATION, + IRVersion=ir_version, + pOutputData=out_ptr, + OutputDataSizeInBytes=size_ptr, + ) + return FeatureDataMLIRExchange720( + MlirInterface=GUID_SUBGRAPH_DECLARATION_REQUEST, + pOutputData=out_ptr, + OutputDataSizeInBytes=size_ptr, + ) + + size = ctypes.c_size_t(0) + first = build(None, ctypes.pointer(size)) + rc = cfs(FEATURE_MLIR_EXCHANGE, ctypes.byref(first), ctypes.sizeof(first)) + if rc < 0: + return None, rc + if size.value == 0: + return b"", 0 + + buf = ctypes.create_string_buffer(size.value) + written = ctypes.c_size_t(size.value) + second = build(ctypes.cast(buf, ctypes.c_void_p), ctypes.pointer(written)) + rc = cfs(FEATURE_MLIR_EXCHANGE, ctypes.byref(second), ctypes.sizeof(second)) + if rc < 0: + return None, rc + if written.value > size.value: + raise click.ClickException( + f"driver overran the response buffer: {written.value} > {size.value}" + ) + return buf.raw[: written.value], 0 + + +#: A line break in any form a Windows producer may emit: CRLF, a run of CRs before LF +#: (CRLF written again through a text-mode stream), or a bare CR. +_LINE_BREAK_RE = re.compile(rb"\r+\n|\r") + + +def normalize_text_payload(data: bytes) -> bytes: + """Normalise a text payload to LF line endings, without its C string terminator. + + The driver writes Windows line endings and reports its size including the NUL + that terminates the string. Both are normalised once, before anything else + reads the payload, so the counts, ``patterns.mlir`` and ``patterns.json`` all + describe exactly the same bytes. Only a text payload is normalised: rewriting + line endings inside MLIR bytecode would corrupt it. + + Normalising is idempotent -- the result contains no CR at all -- so applying it + to already-normalised text is a no-op. + + Args: + data: The text payload as the exchange returned it. + + Returns: + The payload with trailing NULs removed and every line break turned into LF. + """ + # TODO(tests): CRLF, CR-runs before LF, bare CRs and a trailing NUL are not tested + # directly, and nothing asserts normalize(normalize(x)) == normalize(x). + return _LINE_BREAK_RE.sub(b"\n", data.rstrip(b"\x00")) + + +def is_mlir_bytecode(data: bytes) -> bool: + """Report whether *data* is MLIR bytecode rather than text. + + Args: + data: The payload. + + Returns: + True when it carries the bytecode magic (utilities/MlirEncoding.h). + """ + return data.startswith(b"ML\xefR") + + +# ------------------------------------------------------------- counting patterns +# A pattern is one cgc_pattern.pattern declaration. A rule is one flat match +# alternative after any_of expansion: an `any_of { all_of {..} all_of {..} }` block +# contributes one rule per all_of branch, and several blocks in one pattern multiply. +# Do not confuse this with apply_native_constraint "cgc_is_any_of", which is an +# op-family whitelist and no multiplier at all -- searching for the substring any_of +# conflates the two. + +PATTERN_DECL_RE = re.compile(rb"cgc_pattern\.pattern\s+@([A-Za-z0-9_.]+)") +ANY_OF_RE = re.compile(rb"(? str: + """Read a pattern's kind off its body, never its name. + + Args: + body: The pattern body, braces included. + + Returns: + One of :data:`KIND_ORDER`. + """ + if b"jitFunction" in body: + return "kernel" + if b"subgraph_rewrite_desc" in body: + return "fusion" + if b"cgc_add_pattern_cluster_hint" in body: + return "hint" + if b"cgc_mark_pattern_cluster_op" in body: + return "mark" + return "rewrite" + + +def _balanced(text: bytes, open_brace: int) -> int: + """Find the ``}`` closing the ``{`` at *open_brace*. + + Args: + text: The haystack. + open_brace: Index of the opening brace. + + Returns: + Index of the matching close brace, or the last index when unbalanced. + """ + depth = 0 + for token in BRACE_TOKEN_RE.finditer(text, open_brace): + depth += -1 if token.group() == b"}" else 1 + if depth == 0: + return token.start() + return len(text) - 1 + + +def _any_of_branches(body: bytes, open_brace: int) -> int: + """Count ``all_of`` branches directly inside an ``any_of`` block. + + Args: + body: The pattern body. + open_brace: Index of the ``any_of`` block's opening brace. + + Returns: + The branch count, at least 1. + """ + depth = branches = 0 + for token in BRACE_TOKEN_RE.finditer(body, open_brace): + if token.group() == b"}": + depth -= 1 + if depth == 0: + break + else: + depth += 1 + if depth == 2 and token.group() != b"{": # an all_of block + branches += 1 + return branches or 1 + + +def _blank_comments(data: bytes) -> bytes: + """Replace every ``//`` comment outside a string with spaces of the same length. + + Deleting comments would shift every offset after them. Blanking keeps offsets + valid, so a span found here also slices the original bytes -- comments intact, + for display -- and a brace inside a comment can no longer unbalance the match. + + Args: + data: The MLIR text payload. + + Returns: + The same bytes with comment text replaced by spaces. + """ + return STRING_OR_COMMENT_RE.sub( + lambda m: b" " * len(m.group()) if m.group().startswith(b"//") else m.group(), data + ) + + +def _pattern_spans(text: bytes) -> Iterator[tuple[re.Match[bytes], int, int]]: + """Yield each declaration that has a body, with that body's bounds. + + Some patterns declare an ``attributes { depends = .. }`` block between the name and + the body. It is skipped: read as the body, it holds no jitFunction, and a kernel + would be filed under rewrite. + + Args: + text: Comment-blanked MLIR text. + + Yields: + ``(declaration match, open brace, close brace)``. + """ + declarations = list(PATTERN_DECL_RE.finditer(text)) + for index, m in enumerate(declarations): + # The body has to start before the next declaration does. Without that bound a + # declaration with no body of its own adopts the next one's, and both are then + # reported, the first carrying two declarations' text. + limit = declarations[index + 1].start() if index + 1 < len(declarations) else len(text) + try: + i = text.index(b"{", m.end(), limit) + if ATTRIBUTES_TAIL_RE.search(text[m.end() : i]): + i = text.index(b"{", _balanced(text, i) + 1, limit) + except ValueError: + continue # a declaration with no body: nothing to read + yield m, i, _balanced(text, i) + + +def pattern_groups(records: list[dict[str, Any]]) -> list[dict[str, Any]]: + """Bucket records by kind, in :data:`KIND_ORDER`, with each bucket's totals. + + A kind nothing was declared for is left out rather than reported as zero. + + Args: + records: The parsed patterns. + + Returns: + One entry per non-empty kind. + """ + groups: list[dict[str, Any]] = [] + for kind in KIND_ORDER: + members = [r for r in records if r["kind"] == kind] + if not members: + continue + members.sort(key=lambda r: (-(r["benefit"] or 0), r["name"])) + groups.append( + { + "kind": kind, + "members": members, + "patterns": len(members), + "rules": sum(r["rules"] for r in members), + } + ) + return groups + + +def format_pattern_listing(records: list[dict[str, Any]]) -> list[str]: + """Format the ``-v`` listing: every pattern, grouped by what it claims. + + Args: + records: The records from :func:`split_patterns`. + + Returns: + The lines to print. + """ + lines: list[str] = [] + for group in pattern_groups(records): + lines.append( + f"{group['kind']:<8} {group['patterns']:2d} patterns, " + f"{group['rules']:2d} rules {KIND_BLURB[group['kind']]}" + ) + for r in group["members"]: + benefit = " -" if r["benefit"] is None else f"{r['benefit']:4d}" + extra = f" {r['rules']} rules" if r["rules"] > 1 else "" + lines.append(f" {benefit} {r['short_name']}{extra}") + return lines + + +#: A source-file boundary in a merged dump. Anchored on a filename with an +#: extension: a bare "// from" also occurs in prose ("// from outside the root's +#: own chain"), and must not open a section. +SECTION_RE = re.compile(rb"^[ \t]*//[ \t]*from[ \t]+(\S+\.(?:pdll|mlir))[ \t]*$", re.MULTILINE) + +#: The kernel a pattern declares. Driver dumps write ``#cgc.string<"Name">``; the +#: plain string form is accepted as well. +KERNEL_RE = re.compile(rb'jitFunction\s*=\s*(?:#cgc\.string<)?"([^"]+)"') + +#: Identifies the layout of ``patterns.json`` for readers such as the atlas page. +PATTERNS_JSON_FORMAT = "winml-cgc-patterns/1" + + +# TODO(tests): the tests for the prose "// from" anchor and the kernel regex would +# still pass with that guard removed -- they need inputs that fail without it. +def split_patterns(data: bytes) -> list[dict[str, Any]]: + """Split a dump into one record per ``cgc_pattern.pattern`` declaration. + + Each record carries the counts (kind, benefit, rules) and what a reader needs + beyond them: the source file the pattern came from, the kernel it declares, and + the pattern's text with its comments. + + Args: + data: The MLIR text payload, already passed through + :func:`normalize_text_payload`. + + Returns: + Records in declaration order, each with ``index``, ``name``, + ``short_name``, ``kind``, ``benefit``, ``rules``, ``source``, ``kernel``, + ``lines`` and ``mlir``. + """ + text = _blank_comments(data) + sections = [ + (m.start(), m.group(1).decode("utf-8", "replace")) for m in SECTION_RE.finditer(data) + ] + records: list[dict[str, Any]] = [] + for index, (m, open_brace, close_brace) in enumerate(_pattern_spans(text), 1): + head, body = text[m.end() : open_brace], text[open_brace:close_brace] + name = m.group(1).decode("ascii", "replace") + benefit = BENEFIT_RE.search(head) + rules = 1 + for a in ANY_OF_RE.finditer(body): + rules *= _any_of_branches(body, a.end() - 1) + kernel = KERNEL_RE.search(body) + raw = data[m.start() : close_brace + 1] + records.append( + { + "index": index, + "name": name, + "short_name": name.split(".")[-1], # the vendor prefix carries nothing + "kind": _classify(body), + "benefit": int(benefit.group(1)) if benefit else None, + "rules": rules, + "source": next((s for pos, s in reversed(sections) if pos < m.start()), None), + "kernel": kernel.group(1).decode("utf-8", "replace") if kernel else None, + "lines": raw.count(b"\n") + 1, + "mlir": raw.decode("utf-8", "replace"), + } + ) + return records + + +def patterns_document( + data: bytes, adapter: Adapter, sdk: int | None, redist: Path | None +) -> dict[str, Any]: + """Build the ``patterns.json`` document for one dump. + + Args: + data: The MLIR text payload, already passed through + :func:`normalize_text_payload`. + adapter: The adapter the dump describes. + sdk: The redist's SDK version. + redist: The redist directory. + + Returns: + The document: format tag, provenance, summary, per-kind and per-source + grouping, and every pattern. + """ + # TODO(tests): summary.sources, meta.driver_version/redist/encoding and the + # "(unattributed)" source group are not asserted by any test. + patterns = split_patterns(data) + kernels = [p["kernel"] for p in patterns if p["kernel"]] + sources: dict[str, list[int]] = {} + for p in patterns: + sources.setdefault(p["source"] or "(unattributed)", []).append(p["index"]) + return { + "format": PATTERNS_JSON_FORMAT, + "meta": { + "adapter": adapter["description"], + "driver_version": adapter["driver_version"], + "sdk_version": sdk, + "redist": str(redist) if redist else None, + "encoding": "text", + "bytes": len(data), + }, + "summary": { + "patterns": len(patterns), + "rules": sum(p["rules"] for p in patterns), + "kernel_patterns": len(kernels), + "distinct_kernels": len(set(kernels)), + "sources": len(sources), + "top_benefit": max((p["benefit"] or 0 for p in patterns), default=0), + }, + "kinds": {g["kind"]: g["patterns"] for g in pattern_groups(patterns)}, + "sources": [{"name": name, "patterns": ids} for name, ids in sources.items()], + "patterns": patterns, + } + + +def write_patterns_json(path: Path, document: dict[str, Any]) -> None: + """Write ``patterns.json``. + + Args: + path: The file to write. + document: The document from :func:`patterns_document`. + """ + path.write_text( + json.dumps(document, indent=2, ensure_ascii=False) + "\n", encoding="utf-8", newline="\n" + ) + + +def adapter_type(a: Adapter) -> str: + """Classify an adapter for the TYPE column. + + Args: + a: The adapter record. + + Returns: + ``hardware``, ``integrated`` or ``software``. + """ + if not a["is_hardware"]: + return "software" + return "integrated" if a["is_integrated"] else "hardware" + + +def adapter_attrs(a: Adapter) -> str: + """Render the DXCore attributes an adapter advertises. + + Args: + a: The adapter record. + + Returns: + A space-separated subset of ``ML CC GFX``, or ``-``. + """ + return ( + " ".join( + k + for k, v in ( + ("ML", a["generic_ml"]), + ("CC", a["core_compute"]), + ("GFX", a["d3d12_graphics"]), + ) + if v + ) + or "-" + ) + + +def print_adapters(adapters: list[Adapter], stream: Any = None) -> None: + """Print the adapter table. + + Args: + adapters: Every adapter found. + stream: Destination, defaulting to stdout. + """ + out = sys.stdout if stream is None else stream + hdr = f"{'IDX':<3} {'ADAPTER':<30} {'DRIVER':<17} {'TYPE':<10} {'ATTRIBUTES':<11} MLIR" + out.write(hdr + "\n") + out.write("-" * len(hdr) + "\n") + for a in adapters: + mark = {True: "yes", False: "no", None: "?"}[a["mlir"]] + out.write( + f"{a['index']:<3} {a['description'][:30]:<30} {a['driver_version']:<17} " + f"{adapter_type(a):<10} {adapter_attrs(a):<11} {mark}\n" + ) + + +def write_metadata( + path: Path, + adapter: Adapter, + status: str, + encoding: str, + sdk: int | None, + redist: Path | None, + counts: tuple[int, int] | None, + received_bytes: int | None = None, +) -> None: + """Write ``metadata.txt`` beside a dump. + + The first six keys are the ones ``dxcgc-dump-driver-patterns`` writes, unchanged so + existing readers keep working; the rest record which redist and ABI produced the + dump, and how the payload that was stored differs from the one received -- a text + dump is stored with LF line endings. Additive only. + + Args: + path: The file to write. + adapter: The adapter the dump describes. + status: ``dumped``, ``unsupported`` or ``empty``. + encoding: ``text``, ``bytecode`` or ``none``. + sdk: The redist's SDK version. + redist: The redist directory. + counts: ``(patterns, rules)``, when known. + received_bytes: The size of the payload as the exchange returned it, + terminator included, when there was one. + """ + text_file = "patterns.mlir" if status == "dumped" and encoding == "text" else "" + sdk_version = str(sdk) if sdk else "" + patterns, rules = (str(counts[0]), str(counts[1])) if counts else ("", "") + lines: list[tuple[str, str]] = [ + ("adapter", adapter["description"]), + ("driver_version", adapter["driver_version"]), + ("status", status), + ("received_encoding", encoding), + ("text_file", text_file), + ("bytecode_file", "patterns.mlirbc" if encoding == "bytecode" else ""), + ("abi", sdk_version), + ("sdk_version", sdk_version), + ("redist", str(redist) if redist else ""), + ("patterns", patterns), + ("rules", rules), + ("received_bytes", "" if received_bytes is None else str(received_bytes)), + ("line_endings", "lf" if text_file else ""), + ] + path.write_text("".join(f"{k}={v}\n" for k, v in lines), encoding="utf-8", newline="\n") + + +DUMP_FILES = ("patterns.mlir", "patterns.mlirbc", "patterns.json", "metadata.txt") + + +def clear_dump(dest: Path) -> None: + """Remove a previous dump's files from *dest*. + + ``--overwrite`` REPLACES a dump. Without this, re-dumping a driver that now answers + in text leaves the previous run's ``patterns.mlirbc`` beside the new + ``patterns.mlir``, and the directory describes two different dumps at once. + + Args: + dest: The dump directory. + """ + for name in DUMP_FILES: + (dest / name).unlink(missing_ok=True) + + +def dump_dir(adapter: Adapter) -> Path: + """Return the directory a dump for *adapter* is filed under. + + Args: + adapter: The adapter record. + + Returns: + ``patterns//``, relative to the working directory. + """ + # Both parts fall back to a placeholder: an empty one would be joined away, and + # the dump would lose a level of the layout that keeps captures apart. + version = str(adapter["driver_version"]) or "unknown-version" + return Path("patterns") / adapter_slug(adapter["description"]) / version + + +def redist_line(redist: Path | None, sdk: int | None) -> str: + """Render the ``redist:`` line reported on every run. + + Which core answered decides what the rest of the output means. + + Args: + redist: The resolved redist directory, or None. + sdk: Its SDK version, or None. + + Returns: + The line. + """ + if redist is None: + return "redist: none -- MLIR support unknown (run with -v to see where cgc looked)" + return f"redist: {redist} (SDK {sdk}, {host_machine()})" + + +#: The viewer for a patterns.json dump, shipped as package data beside this module. +ATLAS_PAGE = "cgc-pattern-atlas.html" +ATLAS_DIR = "assets" + + +def atlas_page() -> Path: + """Locate the pattern atlas page that ships with this package. + + Asking the package rather than walking the filesystem is what makes ``--open`` + work the same way from a wheel, from an editable install and from a checkout. + + Returns: + The page on disk. + + Raises: + click.ClickException: When the package was installed without its assets. + """ + page = Path(str(resources.files(commands).joinpath(ATLAS_DIR, ATLAS_PAGE))) + if not page.is_file(): + raise click.ClickException( + f"{ATLAS_DIR}/{ATLAS_PAGE} is missing from this install of winml-cli, so " + f"there is no page to open; the dump itself is still written" + ) + return page + + +def read_patterns_json(document: Path) -> dict[str, Any]: + """Read a ``patterns.json``, refusing anything that is not one. + + Every reason a named file cannot be shown is decided here, so ``--open`` reports + a mistyped or unrelated path the way it reports any other bad argument. + + Args: + document: The file named by ``--open``. + + Returns: + The document. + + Raises: + click.UsageError: When it is missing, unreadable, not JSON, or not a + :data:`PATTERNS_JSON_FORMAT` document. + """ + if not document.is_file(): + raise click.UsageError(f"--open: {document} is not a file") + try: + raw = json.loads(document.read_text(encoding="utf-8")) + except (OSError, UnicodeDecodeError, json.JSONDecodeError) as e: + raise click.UsageError(f"--open: {document} could not be read as JSON: {e}") from e + # Anything that is not a mapping fails the format check below, and "meta" has to + # be one too: open_atlas labels the page through it. + data: dict[str, Any] = cast("dict[str, Any]", raw) if isinstance(raw, dict) else {} + if data.get("format") != PATTERNS_JSON_FORMAT or not isinstance(data.get("meta", {}), dict): + raise click.UsageError( + f"--open: {document} is not a {PATTERNS_JSON_FORMAT} document; pass a " + f"patterns.json written by --dump" + ) + return data + + +def open_atlas(document: Path | None) -> None: + """Open the pattern atlas in a browser, showing *document* when there is one. + + The page reads a file the viewer picks, and a page opened from disk may not read + one beside it, so a dump is embedded into a throwaway copy of the page instead. + + Args: + document: A ``patterns.json`` to show, or None to open the empty page. + + Raises: + click.UsageError: When *document* is not a patterns.json. + """ + page = atlas_page() + if document is not None: + data = read_patterns_json(document) + data.setdefault("meta", {})["label"] = str(document) + # Every "<" is escaped, so no pattern's own text can close the element early. + payload = json.dumps(data).replace("<", "\\u003c") + seeded = page.read_text(encoding="utf-8").replace( + '', + f'', + ) + # Named after the dump, not a fresh temp file per run: the browser reads the + # page after this process is gone, so it cannot be cleaned up on the way out, + # and two dumps opened together must not land on the same name. + stem = adapter_slug(f"{document.parent.parent.name}-{document.parent.name}") + page = Path(tempfile.gettempdir()) / f"winml-{stem}-{ATLAS_PAGE}" + page.write_text(seeded, encoding="utf-8", newline="\n") + click.echo(f"opening {page}") + webbrowser.open(page.as_uri()) + + +@dataclass(kw_only=True) +class Options: + """The resolved flags for one ``adapter`` invocation. + + Attributes: + adapter: The ``-a`` selector. + verbose: The ``-v`` flag. + d3d12_dir: The ``--d3d12-dir`` value. + dump: The ``--dump`` flag. + overwrite: The ``--overwrite`` flag. + open_atlas: The ``--open`` value: None when it was not passed, an empty + string for this run's own dump, else the ``patterns.json`` to show. + """ + + adapter: str = "" + verbose: bool = False + d3d12_dir: str = "" + dump: bool = False + overwrite: bool = False + open_atlas: str | None = None + + +# TODO(tests): nothing drives report_patterns end to end: the dumped, unsupported +# and empty branches, the patterns.json write, --overwrite, and the refusal. +def report_patterns(opts: Options, adapter: Adapter, redist: Path, sdk: int) -> Path | None: + """Report, and optionally dump, what one adapter's driver declares. + + Args: + opts: The flags. + adapter: The selected adapter. + redist: The resolved redist. + sdk: Its SDK version. + + Returns: + The ``patterns.json`` this run wrote, for :func:`open_atlas`, or None. + """ + dest = dump_dir(adapter) + if opts.dump: + # An empty directory, left by a run that died between mkdir and the writes, is + # reused; guard_output only refuses one that holds files. + cli_utils.guard_output(dest, opts.overwrite, label="Dump") + + desc = adapter["description"] + status, encoding, received, data, doc = "unsupported", "none", None, b"", None + device = probe_mlir_support(adapter, redist, sdk, verbose=opts.verbose) + if device is None: + click.echo(f"{desc}: driver does not implement MLIR programs") + else: + try: + payload, rc = mlir_exchange(device, sdk, adapter["ir_version"]) + finally: + release(device) + if payload is None: + # The driver said it implements the exchange and then refused it. + raise click.ClickException( + f"MLIR exchange failed after the driver claimed support: {hrs(rc)}" + ) + received = len(payload) + bytecode = is_mlir_bytecode(payload) + # Normalised once, here, so a payload that is only its terminator is empty. + data = payload if bytecode else normalize_text_payload(payload) + if not data: + status = "empty" + click.echo(f"{desc}: driver returned an empty subgraph declaration") + elif bytecode: + status, encoding = "dumped", "bytecode" + click.echo( + f"{desc}: {len(data)} bytes of MLIR bytecode; rendering it " + f"as text needs cgc-opt, so cgc cannot count it" + ) + else: + status, encoding = "dumped", "text" + doc = patterns_document(data, adapter, sdk, redist) + click.echo( + f"{desc}: {doc['summary']['patterns']} patterns / {doc['summary']['rules']} rules" + ) + if opts.verbose: + click.echo("") + for line in format_pattern_listing(doc["patterns"]): + click.echo(line) + + if not opts.dump: + return None + dest.mkdir(parents=True, exist_ok=True) + clear_dump(dest) + if status == "dumped": + name = "patterns.mlir" if encoding == "text" else "patterns.mlirbc" + (dest / name).write_bytes(data) + click.echo(f"wrote {dest / name} ({len(data)} bytes)") + counts = None + if doc is not None: + # Bytecode cannot be split without the real MLIR parser, so only a text dump + # gets the structured copy. + write_patterns_json(dest / "patterns.json", doc) + click.echo(f"wrote {dest / 'patterns.json'}") + counts = (doc["summary"]["patterns"], doc["summary"]["rules"]) + write_metadata(dest / "metadata.txt", adapter, status, encoding, sdk, redist, counts, received) + click.echo(f"wrote {dest / 'metadata.txt'}") + return dest / "patterns.json" if doc is not None else None + + +# TODO(tests): the Click exceptions raised on these paths are untested as a caller +# sees them -- exit 1 for "could not ask", exit 2 for bad arguments. +def resolve_run(opts: Options) -> tuple[Path | None, int | None, tuple[str, Path] | None]: + """Resolve the redist both subcommands answer through, and report which it is. + + Args: + opts: The flags. + + Returns: + ``(redist, sdk_version, named)``; *named* is the source the caller named it + by, when they named one. + + Raises: + click.ClickException: When a named redist cannot be used. + """ + redist, sdk, tried = resolve_redist(opts.d3d12_dir) + named = _named_redist(opts.d3d12_dir) + if redist is None and (named or opts.verbose): + # Asked for a specific redist, or asked to be told: an unusable one is fatal. + click.echo(redist_failure(tried), err=True) + if named: + raise click.ClickException( + f"{named[0]} named a redist that cannot be used; refusing to answer " + f"about a different one." + ) + # Which core answered decides what every other line means, so it is always + # reported. When there is none the line says so, and the MLIR column reads ? + # rather than no: the question cannot be asked, so it is left unanswered. + click.echo(redist_line(redist, sdk)) + return redist, sdk, named + + +# TODO(tests): the Click exceptions raised on these paths are untested as a caller +# sees them -- exit 1 for "could not ask", exit 2 for bad arguments. +def run_adapters(opts: Options) -> None: + """Run the ``adapters`` subcommand: the table of every D3D12 adapter. + + Args: + opts: The flags. + """ + adapters = list_adapters() + if not adapters: + click.echo("no D3D12 adapters found") + return + + try: + redist, sdk, named = resolve_run(opts) + if redist is not None: + try: + for a in adapters: + release(probe_mlir_support(a, redist, sdk or 0, verbose=opts.verbose)) + except RedistUnusable as exc: + # D3D12 cannot use this redist at all. One the caller named must be + # usable, so that fails; otherwise the listing stands, MLIR unanswered. + if named: + raise + for a in adapters: + a["mlir"] = None + click.echo(f"cgc: {exc.format_message()}", err=True) + print_adapters(adapters) + finally: + close_adapters(adapters) + + +def run_patterns(opts: Options) -> None: + """Run the ``patterns`` subcommand: what a driver declares it can match. + + With no ``-a`` every adapter is asked, so a machine's answer is one command. + + Args: + opts: The flags. + + Raises: + click.UsageError: When ``--dump`` did not name its adapter, or ``--open`` + named something that is not a patterns.json. + click.ClickException: When an adapter could not be asked. Sweeping every + adapter, the others are still reported first. + """ + if opts.dump and opts.open_atlas: + # Two dumps to show, this run's and a named one. Silently preferring either + # throws away what the caller asked for. + raise click.UsageError( + "--open takes a patterns.json only without --dump; with --dump it shows " + "the dump this run writes" + ) + if opts.open_atlas is not None and not opts.dump: + # Nothing of this run's own to show: a named dump, or the empty page to drop + # a file on. Either way no driver is asked. + open_atlas(Path(opts.open_atlas) if opts.open_atlas else None) + return + if opts.dump and not opts.adapter: + # Required rather than defaulted: a dump is filed under a slug and a driver + # version, and dumping whichever adapter happened to be first writes a + # correct-looking directory for the wrong device. + raise click.UsageError("--dump needs an adapter: name one with -a ") + adapters = list_adapters() + if not adapters: + click.echo("no D3D12 adapters found") + return + + try: + redist, sdk, _ = resolve_run(opts) + if redist is None or sdk is None: + # One condition for the whole run: nothing was asked, so no adapter is + # named as having failed. + raise click.ClickException( + "cgc patterns needs an Agility SDK redist; none was found " + "(run with -v to see where cgc looked)" + ) + targets = [select_adapter(adapters, opts.adapter)] if opts.adapter else adapters + written: Path | None = None + unanswered: list[str] = [] + for adapter in targets: + try: + written = report_patterns(opts, adapter, redist, sdk) or written + except RedistUnusable: # noqa: PERF203 - next to a device, try costs nothing + raise # no adapter can be asked through this redist + except click.ClickException as exc: + # Sweeping every adapter: one driver that cannot answer must not cost + # the others their line. The run still fails, once, at the end. + if len(targets) == 1: + raise + click.echo(f"cgc: {adapter['description']}: {exc.format_message()}", err=True) + unanswered.append(adapter["description"]) + finally: + close_adapters(adapters) + if opts.open_atlas is not None: + if written is None: + click.echo("cgc: no patterns.json was written, so the atlas opens empty", err=True) + open_atlas(written) + if unanswered: + raise click.ClickException(f"could not ask: {', '.join(unanswered)}") + + +# A plain Group: the root LazyGroup already instruments `cgc` itself, and an +# ActionGroup here would instrument the subcommand as well, so every run would count +# twice where every other command counts once. +@click.group( + name="cgc", + invoke_without_command=True, +) +@click.pass_context +def cgc(ctx: click.Context) -> None: + """Inspect D3D12 adapters and the MLIR patterns their drivers declare. + + Answers two questions on any Windows machine: which D3D12 adapters are here + (`adapters`), and what patterns does a driver declare it can match (`patterns`). + + Querying MLIR support or patterns needs an Agility SDK D3D12Core.dll of SDK 720 or + newer; point at one with --d3d12-dir or $WINML_D3D12_DIR. Listing adapters needs + neither a redist nor a build. `patterns --open` without --dump only opens the + viewer or an existing dump and needs neither a GPU nor a redist. + """ + # Mirror the root `winml` group: a bare invocation prints help and succeeds, + # rather than Click's default usage error. + if ctx.invoked_subcommand is None: + click.echo(ctx.get_help()) + ctx.exit(0) + + +@cgc.command("adapters") +@click.option( + "--d3d12-dir", + default="", + metavar="PATH", + help="Directory holding D3D12Core.dll ($WINML_D3D12_DIR, $D3D12_DIR).", +) +@cli_utils.verbosity_options() +@click.pass_context +def adapters_cmd(ctx: click.Context, d3d12_dir: str, verbose: int, quiet: bool) -> None: + """List DXCore ML or core-compute adapters and their MLIR support. + + Lists generic-ML adapters, or core-compute adapters if the first list is empty. + Needs nothing beyond Windows. The MLIR column is answered through the Agility + SDK redist named on the redist: line -- without one it reads ?, since the question + cannot be asked. `winml cgc patterns` then reports what a driver declares. + + \b + Examples: + # Every adapter, with its driver version and MLIR support + winml cgc adapters + + \b + # Show the raw capability HRESULT behind each answer + winml cgc adapters -v + + \b + # Ask through a particular Agility SDK redist + winml cgc adapters --d3d12-dir + """ # noqa: D301 - the \b escapes must become backspace characters + verbose, quiet = cli_utils.resolve_verbosity(ctx, verbose, quiet) + configure_logging(verbosity=verbose, quiet=quiet) + opts = Options(verbose=bool(verbose), d3d12_dir=d3d12_dir) + logger.debug("cgc adapters: %s", vars(opts)) + run_adapters(opts) + + +@cgc.command("patterns") +@click.option( + "-a", + "--adapter", + "adapter_sel", + default="", + metavar="SUBSTRING|INDEX", + help="The adapter to ask, by description substring or index; --dump needs it. " + "Without it, every adapter is asked.", +) +@click.option( + "--d3d12-dir", + default="", + metavar="PATH", + help="Directory holding D3D12Core.dll ($WINML_D3D12_DIR, $D3D12_DIR).", +) +@click.option( + "--dump", + is_flag=True, + default=False, + help="Write a dump under ./patterns (needs -a). Non-empty text produces patterns.mlir and " + "patterns.json; bytecode produces patterns.mlirbc. Successful queries also write metadata.txt.", +) +@cli_utils.overwrite_option() +@click.option( + "--open", + "open_atlas", + is_flag=False, + flag_value="", + default=None, + metavar="[PATTERNS.JSON]", + help="Open the pattern atlas. Without --dump, show a named patterns.json or an empty " + "page without querying drivers. With --dump, show this run's dump when available.", +) +@cli_utils.verbosity_options() +@click.pass_context +def patterns_cmd( + ctx: click.Context, + adapter_sel: str, + d3d12_dir: str, + dump: bool, + overwrite: bool, + open_atlas: str | None, + verbose: int, + quiet: bool, +) -> None: + """Report the MLIR patterns a driver declares it can match. + + Querying a driver needs an Agility SDK D3D12Core.dll of SDK 720 or newer: name one + with --d3d12-dir, or set $WINML_D3D12_DIR. With no -a every enumerated adapter is + asked. --dump writes under the adapter description and driver version; adapters + with the same normalized description share a destination. + + --open without --dump opens an existing JSON dump or an empty atlas without + querying drivers, so it needs neither a GPU nor a redist. + + stdout carries results and stderr diagnostics. A multi-adapter run can print + partial results before exiting 1. Exit 0 includes unsupported or empty answers; + exit 1 means the operation could not complete; exit 2 is bad arguments. + + \b + Examples: + # What every adapter's driver declares + winml cgc patterns + + \b + # One adapter, by substring or index + winml cgc patterns -a nvidia + + \b + # Every pattern it declares, grouped by kind + winml cgc patterns -a nvidia -v + + \b + # Keep the declaration on disk, and look at it + winml cgc patterns -a nvidia --dump --open + + \b + # Look at a dump taken earlier, with no hardware involved + winml cgc patterns --open patterns/nvidia.../32.0.16.3004/patterns.json + + \b + # Ask a different redist the same question + winml cgc patterns -a nvidia --d3d12-dir + """ # noqa: D301 - the \b escapes must become backspace characters + # Merge top-level -v/-q with subcommand-level flags so either position works. + verbose, quiet = cli_utils.resolve_verbosity(ctx, verbose, quiet) + + # Standard verbosity contract: stderr-only logs in the shared format. `-v` + # here keeps its cgc-specific second job of adding the raw capability + # HRESULTs and the full pattern listing. + configure_logging(verbosity=verbose, quiet=quiet) + + opts = Options( + adapter=adapter_sel, + verbose=bool(verbose), + d3d12_dir=d3d12_dir, + dump=dump, + overwrite=overwrite, + open_atlas=open_atlas, + ) + logger.debug("cgc patterns: %s", vars(opts)) + run_patterns(opts) diff --git a/tests/unit/commands/test_cgc.py b/tests/unit/commands/test_cgc.py new file mode 100644 index 000000000..a40ed16f7 --- /dev/null +++ b/tests/unit/commands/test_cgc.py @@ -0,0 +1,1299 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Tests for the hardware-free half of ``winml cgc adapters`` and ``cgc patterns``. + +Every expectation here is derived from data this module builds, never from a +transcript: MLIR sources are assembled from a known number of patterns and +``any_of`` branches and the counts are computed from those parameters, and the +PE files are synthesized around a chosen ``D3D12SDKVersion`` value which is then +read back. The DXCore/D3D12 half needs real hardware and is not covered here. +""" + +from __future__ import annotations + +import ctypes +import io +import json +import re +import struct +from pathlib import Path +from urllib.parse import urlparse +from urllib.request import url2pathname + +import click +import pytest + +from winml.modelkit.commands.cgc import ( + COMMENT_RE, + DUMP_FILES, + MIN_SDK_VERSION, + PATTERNS_JSON_FORMAT, + REDIST_ENV_VARS, + FeatureDataMLIRExchange721, + Options, + RedistUnusable, + _named_redist, + adapter_attrs, + adapter_slug, + adapter_type, + atlas_page, + clear_dump, + close_adapters, + dump_dir, + format_driver_version, + format_pattern_listing, + hrs, + inspect_redist, + is_mlir_bytecode, + list_adapters, + make_version_number, + normalize_text_payload, + open_atlas, + pattern_groups, + patterns_document, + pe_export_u32, + print_adapters, + probe_mlir_support, + redist_candidates, + redist_failure, + redist_line, + release, + resolve_redist, + run_patterns, + select_adapter, + split_patterns, + write_metadata, + write_patterns_json, +) + + +#: The documented exit codes: 1 when the question could not be asked, 2 for bad +#: arguments. +EXIT_FAIL, EXIT_USAGE = 1, 2 + + +def path_of(uri: str) -> Path: + """Turn a file:// URI back into a path, undoing as_uri()'s percent-encoding.""" + return Path(url2pathname(urlparse(uri).path)) + + +def count_patterns(data: bytes) -> tuple[int, int]: + """Count patterns and rules the way the command reports them.""" + records = split_patterns(data) + return len(records), sum(r["rules"] for r in records) + + +# --------------------------------------------------------------- MLIR builders + +#: One marker per kind, taken from the classifier's own contract: a body holding +#: this token must be filed under that kind. +KIND_MARKERS = { + "kernel": "jitFunction", + "fusion": "subgraph_rewrite_desc", + "hint": "cgc_add_pattern_cluster_hint", + "mark": "cgc_mark_pattern_cluster_op", + "rewrite": "some_other_op", +} + + +def make_pattern(name: str, kind: str, benefit: int | None = None, branches: int = 1) -> str: + """Build one ``cgc_pattern.pattern`` whose kind and rule count are known. + + Args: + name: Pattern name. + kind: One of :data:`KIND_MARKERS`. + benefit: Optional benefit value. + branches: Number of ``all_of`` alternatives inside one ``any_of`` block; + 1 emits no ``any_of`` at all, so the pattern expands to a single rule. + + Returns: + The pattern source. + """ + head = f"cgc_pattern.pattern @{name}" + if benefit is not None: + head += f" benefit({benefit})" + body = [f" {KIND_MARKERS[kind]} : i32"] + if branches > 1: + alts = "\n".join(f" all_of {{ op_{i} }}" for i in range(branches)) + body.append(" any_of {\n" + alts + "\n }") + return head + " {\n" + "\n".join(body) + "\n}\n" + + +# ----------------------------------------------------------------- PE builders + +_MACHINE_FOR_ARCH = {"x86": 0x014C, "x64": 0x8664, "arm64": 0xAA64} + + +def make_pe( + sdk_version: int | None, arch: str = "x64", export_name: bytes = b"D3D12SDKVersion" +) -> bytes: + """Synthesize a minimal PE32+ exporting one UINT32 data symbol. + + The single section maps RVA to file offset one-to-one, so the export walk has + a real table to traverse rather than a stub. + + Args: + sdk_version: Value the export should carry, or None to emit no export + directory at all. + arch: Machine to stamp into the COFF header. + export_name: Name to publish the value under. + + Returns: + The PE image bytes. + """ + pe_off, sec_rva = 0x80, 0x400 + buf = bytearray(0x1400) + buf[0:2] = b"MZ" + struct.pack_into(" no export table + + struct.pack_into(" Path: + """Write a synthesized ``D3D12Core.dll`` into *directory*. + + Args: + directory: Destination, created if absent. + sdk_version: Value for the ``D3D12SDKVersion`` export. + arch: Machine to stamp. + + Returns: + The directory. + """ + directory.mkdir(parents=True, exist_ok=True) + (directory / "D3D12Core.dll").write_bytes(make_pe(sdk_version, arch)) + return directory + + +# ------------------------------------------------------------------- adapters + + +def make_adapter(index: int, description: str, *, hardware: bool, integrated: bool = False) -> dict: + """Build an adapter record of the shape ``list_adapters`` returns. + + Args: + index: Selection index. + description: DXCore DriverDescription. + hardware: Whether the adapter is hardware. + integrated: Whether it is integrated. + + Returns: + The record. + """ + return { + "index": index, + "description": description, + "driver_version": "1.2.3.4", + "is_hardware": hardware, + "is_integrated": integrated, + "generic_ml": True, + "core_compute": True, + "d3d12_graphics": True, + "mlir": None, + } + + +class TestAdapterSlug: + """The dump directory name derived from an adapter description.""" + + @pytest.mark.parametrize( + ("description", "expected"), + [ + ("NVIDIA GeForce RTX 5090 D", "nvidia-geforce-rtx-5090-d"), + ("AMD Radeon(TM) 880M Graphics", "amd-radeon-880m-graphics"), + ("Intel(R) Graphics", "intel-graphics"), + ("Microsoft Basic Render Driver", "microsoft-basic-render-driver"), + ], + ) + def test_known_descriptions(self, description: str, expected: str) -> None: + assert adapter_slug(description) == expected + + def test_leading_and_trailing_separators_are_trimmed(self) -> None: + assert adapter_slug(" Weird Name!!! ") == "weird-name" + + def test_slug_is_filesystem_safe(self) -> None: + slug = adapter_slug(r'A/B\C:D*E?F"GI|J') + assert not set(slug) & set('/\\:*?"<>|') + + +class TestDriverVersion: + """The DXCore u64 rendered as four 16-bit parts.""" + + @pytest.mark.parametrize( + "parts", + [(32, 0, 16, 3004), (0, 0, 0, 0), (0xFFFF, 0xFFFF, 0xFFFF, 0xFFFF), (10, 0, 26100, 9278)], + ) + def test_round_trip_through_make_version_number(self, parts: tuple[int, int, int, int]) -> None: + packed = make_version_number(*parts) + assert format_driver_version(packed) == ".".join(str(p) for p in parts) + + def test_high_word_first(self) -> None: + assert format_driver_version(1 << 48) == "1.0.0.0" + + +class TestHresultRendering: + """HRESULTs print as hex, with a name when one is known.""" + + def test_known_code_carries_its_name(self) -> None: + assert hrs(0x80070057) == "0x80070057 (E_INVALIDARG)" + + def test_unknown_code_prints_raw(self) -> None: + assert hrs(0x81234567) == "0x81234567" + + def test_negative_input_is_rendered_unsigned(self) -> None: + assert hrs(-2147024809) == hrs(0x80070057) + + +class TestPatternParser: + """Counting patterns and the rules they expand to.""" + + def test_pattern_count_matches_what_was_built(self) -> None: + built = 7 + src = "".join(make_pattern(f"p.n{i}", "kernel") for i in range(built)) + patterns, rules = count_patterns(src.encode()) + assert (patterns, rules) == (built, built) + + def test_any_of_branches_multiply_into_rules(self) -> None: + branches = [1, 2, 3, 4] + src = "".join(make_pattern(f"p.n{i}", "kernel", branches=b) for i, b in enumerate(branches)) + patterns, rules = count_patterns(src.encode()) + assert patterns == len(branches) + assert rules == sum(branches) + + def test_comments_are_stripped_before_parsing(self) -> None: + src = "// cgc_pattern.pattern @commented.out {\n" + make_pattern("p.real", "kernel") + assert count_patterns(src.encode()) == (1, 1) + + def test_is_any_of_constraint_is_not_a_rule_multiplier(self) -> None: + plain = make_pattern("p.plain", "kernel") + decoyed = plain.replace( + " jitFunction : i32", + ' jitFunction : i32\n apply_native_constraint "cgc_is_any_of"', + ) + assert count_patterns(decoyed.encode()) == count_patterns(plain.encode()) + + def test_attributes_block_is_not_mistaken_for_the_body(self) -> None: + src = ( + "cgc_pattern.pattern @p.deps attributes { depends = [@other] } {\n" + " jitFunction : i32\n}\n" + ) + assert split_patterns(src.encode())[0]["kind"] == "kernel" + + def test_declaration_without_a_body_is_skipped(self) -> None: + assert count_patterns(b"cgc_pattern.pattern @p.nobody") == (0, 0) + + @pytest.mark.parametrize("kind", sorted(KIND_MARKERS)) + def test_kind_is_read_from_the_body_not_the_name(self, kind: str) -> None: + # Name every pattern after a *different* kind to prove the name is ignored. + misleading = "kernel" if kind != "kernel" else "rewrite" + src = make_pattern(f"vendor.{misleading}_looking_name", kind) + assert split_patterns(src.encode())[0]["kind"] == kind + + def test_groups_are_ordered_and_totalled(self) -> None: + counts = {"kernel": 3, "fusion": 2, "mark": 4} + src = "".join( + make_pattern(f"p.{kind}{i}", kind) for kind, n in counts.items() for i in range(n) + ) + groups = pattern_groups(split_patterns(src.encode())) + assert [g["kind"] for g in groups] == ["kernel", "fusion", "mark"] + assert {g["kind"]: g["patterns"] for g in groups} == counts + + def test_empty_kinds_are_omitted(self) -> None: + groups = pattern_groups(split_patterns(make_pattern("p.only", "hint").encode())) + assert [g["kind"] for g in groups] == ["hint"] + + def test_members_sort_by_descending_benefit(self) -> None: + src = "".join(make_pattern(f"p.b{b}", "kernel", benefit=b) for b in (10, 200, 36)) + members = pattern_groups(split_patterns(src.encode()))[0]["members"] + assert [m["benefit"] for m in members] == [200, 36, 10] + + +class TestPeExportParser: + """Reading D3D12SDKVersion without loading the DLL.""" + + @pytest.mark.parametrize("version", [614, 616, 720, 721, 1]) + def test_reads_back_the_value_it_was_given(self, tmp_path: Path, version: int) -> None: + dll = tmp_path / "D3D12Core.dll" + dll.write_bytes(make_pe(version)) + assert pe_export_u32(dll, b"D3D12SDKVersion") == ("x64", version) + + def test_reports_machine_for_every_arch(self, tmp_path: Path) -> None: + for arch in _MACHINE_FOR_ARCH: + dll = tmp_path / f"{arch}.dll" + dll.write_bytes(make_pe(720, arch)) + assert pe_export_u32(dll, b"D3D12SDKVersion")[0] == arch + + def test_missing_export_yields_no_value(self, tmp_path: Path) -> None: + dll = tmp_path / "D3D12Core.dll" + dll.write_bytes(make_pe(720, export_name=b"SomethingElse")) + assert pe_export_u32(dll, b"D3D12SDKVersion") == ("x64", None) + + def test_absent_export_directory_yields_no_value(self, tmp_path: Path) -> None: + dll = tmp_path / "D3D12Core.dll" + dll.write_bytes(make_pe(None)) + assert pe_export_u32(dll, b"D3D12SDKVersion") == ("x64", None) + + def test_non_pe_file_is_rejected(self, tmp_path: Path) -> None: + dll = tmp_path / "D3D12Core.dll" + dll.write_bytes(b"not a PE at all" + bytes(64)) + assert pe_export_u32(dll, b"D3D12SDKVersion") == (None, None) + + +class TestInspectRedist: + """Deciding whether a directory is a usable Agility SDK redist.""" + + def test_accepts_a_core_at_the_floor(self, tmp_path: Path) -> None: + version, reason = inspect_redist(write_redist(tmp_path / "ok", MIN_SDK_VERSION)) + assert version == MIN_SDK_VERSION + assert reason == f"SDK {MIN_SDK_VERSION}" + + def test_accepts_a_newer_core(self, tmp_path: Path) -> None: + assert inspect_redist(write_redist(tmp_path / "new", MIN_SDK_VERSION + 1))[0] is not None + + def test_rejects_a_core_below_the_floor(self, tmp_path: Path) -> None: + version, reason = inspect_redist(write_redist(tmp_path / "old", MIN_SDK_VERSION - 1)) + assert version is None + assert str(MIN_SDK_VERSION) in reason + + def test_rejects_a_missing_directory(self, tmp_path: Path) -> None: + assert inspect_redist(tmp_path / "absent") == (None, "missing") + + def test_rejects_an_empty_directory(self, tmp_path: Path) -> None: + empty = tmp_path / "empty" + empty.mkdir() + assert inspect_redist(empty) == (None, "empty") + + def test_rejects_a_directory_without_the_core(self, tmp_path: Path) -> None: + other = tmp_path / "other" + other.mkdir() + (other / "readme.txt").write_text("nothing here", encoding="utf-8") + assert inspect_redist(other) == (None, "no D3D12Core.dll") + + def test_rejects_a_foreign_architecture(self, tmp_path: Path) -> None: + from winml.modelkit.commands.cgc import host_machine + + foreign = next(a for a in _MACHINE_FOR_ARCH if a != host_machine()) + version, reason = inspect_redist(write_redist(tmp_path / "arch", 720, foreign)) + assert version is None + assert foreign in reason + + +class TestRedistCandidates: + """Which directories are searched, and in what order.""" + + def test_explicit_flag_is_the_only_candidate(self, monkeypatch: pytest.MonkeyPatch) -> None: + for var in REDIST_ENV_VARS: + monkeypatch.setenv(var, r"C:\from-env") + candidates = list(redist_candidates(r"C:\from-flag")) + assert candidates == [("--d3d12-dir", Path(r"C:\from-flag"))] + + def test_first_environment_variable_set_wins_and_stops( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + for var in REDIST_ENV_VARS: + monkeypatch.setenv(var, f"C:\\{var}") + candidates = list(redist_candidates()) + assert candidates == [(f"${REDIST_ENV_VARS[0]}", Path(f"C:\\{REDIST_ENV_VARS[0]}"))] + + def test_later_variable_is_used_when_the_first_is_unset( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.delenv(REDIST_ENV_VARS[0], raising=False) + monkeypatch.setenv(REDIST_ENV_VARS[1], r"C:\second") + assert next(iter(redist_candidates()))[0] == f"${REDIST_ENV_VARS[1]}" + + def test_bin_roots_are_probed_when_nothing_is_named( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + for var in REDIST_ENV_VARS: + monkeypatch.delenv(var, raising=False) + labels = [label for label, _ in redist_candidates()] + assert labels, "expected implicit bin/ roots" + assert all("bin" in label for label in labels) + + def test_each_root_is_probed_as_d3d12_then_bare_bin( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + for var in REDIST_ENV_VARS: + monkeypatch.delenv(var, raising=False) + paths = [p for _, p in redist_candidates()] + for d3d12 in (p for p in paths if p.name == "D3D12"): + assert d3d12.parent in paths, f"bare bin/ missing for {d3d12}" + + def test_candidates_are_not_repeated(self, monkeypatch: pytest.MonkeyPatch) -> None: + for var in REDIST_ENV_VARS: + monkeypatch.delenv(var, raising=False) + paths = [p for _, p in redist_candidates()] + assert len(paths) == len(set(paths)) + + +class TestNamedRedist: + """Whether the caller named a redist by hand, and by which source.""" + + def test_flag_is_reported(self, monkeypatch: pytest.MonkeyPatch) -> None: + for var in REDIST_ENV_VARS: + monkeypatch.delenv(var, raising=False) + assert _named_redist(r"C:\x") == ("--d3d12-dir", Path(r"C:\x")) + + def test_flag_outranks_the_environment(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv(REDIST_ENV_VARS[0], r"C:\env") + assert _named_redist(r"C:\x") == ("--d3d12-dir", Path(r"C:\x")) + + def test_environment_is_reported(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv(REDIST_ENV_VARS[0], r"C:\env") + assert _named_redist() == (f"${REDIST_ENV_VARS[0]}", Path(r"C:\env")) + + def test_nothing_named_returns_none(self, monkeypatch: pytest.MonkeyPatch) -> None: + for var in REDIST_ENV_VARS: + monkeypatch.delenv(var, raising=False) + assert _named_redist() is None + + +class TestSelectAdapter: + """Resolving ``-a`` to one adapter.""" + + @pytest.fixture + def adapters(self) -> list[dict]: + return [ + make_adapter(0, "Microsoft Basic Render Driver", hardware=False), + make_adapter(1, "NVIDIA GeForce RTX 5090 D", hardware=True), + make_adapter(2, "AMD Radeon(TM) 880M Graphics", hardware=True, integrated=True), + ] + + def test_default_prefers_the_first_hardware_adapter(self, adapters: list[dict]) -> None: + assert select_adapter(adapters, "")["index"] == 1 + + def test_default_falls_back_to_index_zero(self) -> None: + software = [make_adapter(0, "Software", hardware=False)] + assert select_adapter(software, "")["index"] == 0 + + def test_bare_index_selects_positionally(self, adapters: list[dict]) -> None: + assert select_adapter(adapters, "2")["index"] == 2 + + def test_substring_is_case_insensitive(self, adapters: list[dict]) -> None: + assert select_adapter(adapters, "nvidia")["index"] == 1 + + def test_a_number_that_is_no_index_is_matched_as_a_substring( + self, adapters: list[dict] + ) -> None: + # "5090" is a model number, not an index, on a three-adapter machine. + assert select_adapter(adapters, "5090")["index"] == 1 + + def test_an_index_wins_over_a_substring(self, adapters: list[dict]) -> None: + # A number that names a real index stays positional, and is never retried + # as a substring. + assert select_adapter(adapters, "2")["index"] == 2 + + def test_out_of_range_index_is_a_usage_error(self, adapters: list[dict]) -> None: + with pytest.raises(click.UsageError) as excinfo: + select_adapter(adapters, "99") + assert excinfo.value.exit_code == EXIT_USAGE + + def test_unmatched_substring_is_a_usage_error(self, adapters: list[dict]) -> None: + with pytest.raises(click.UsageError) as excinfo: + select_adapter(adapters, "nosuchvendor") + assert excinfo.value.exit_code == EXIT_USAGE + + +class TestFlagRules: + """How the flags constrain one another.""" + + def test_dump_without_an_adapter_is_a_usage_error(self) -> None: + # Refused before DXCore is touched, so this needs no hardware. + with pytest.raises(click.UsageError) as excinfo: + run_patterns(Options(dump=True)) + assert excinfo.value.exit_code == EXIT_USAGE + + def test_click_exceptions_carry_the_documented_codes(self) -> None: + # The command raises these rather than exiting itself, so the documented + # contract -- 1 could not ask, 2 bad arguments -- rests on Click's values. + assert click.UsageError("x").exit_code == EXIT_USAGE + assert click.ClickException("x").exit_code == EXIT_FAIL + + +class TestWriteMetadata: + """The provenance file written beside a dump.""" + + @staticmethod + def read_keys(path: Path) -> dict[str, str]: + text = path.read_text(encoding="utf-8") + return dict(line.split("=", 1) for line in text.splitlines()) + + def test_records_counts_for_a_successful_dump(self, tmp_path: Path) -> None: + target = tmp_path / "metadata.txt" + adapter = make_adapter(0, "NVIDIA GeForce RTX 5090 D", hardware=True) + write_metadata(target, adapter, "dumped", "text", 720, Path(r"C:\redist"), (51, 63)) + keys = self.read_keys(target) + assert keys["status"] == "dumped" + assert keys["patterns"] == "51" + assert keys["rules"] == "63" + assert keys["text_file"] == "patterns.mlir" + assert keys["bytecode_file"] == "" + assert keys["driver_version"] == adapter["driver_version"] + + def test_leaves_counts_empty_when_unavailable(self, tmp_path: Path) -> None: + target = tmp_path / "metadata.txt" + adapter = make_adapter(0, "AMD Radeon(TM) 880M Graphics", hardware=True) + write_metadata(target, adapter, "unsupported", "none", 720, Path(r"C:\redist"), None) + keys = self.read_keys(target) + assert keys["status"] == "unsupported" + assert keys["patterns"] == "" + assert keys["rules"] == "" + assert keys["text_file"] == "" + + def test_bytecode_sets_only_the_bytecode_file(self, tmp_path: Path) -> None: + target = tmp_path / "metadata.txt" + adapter = make_adapter(0, "Vendor Device", hardware=True) + write_metadata(target, adapter, "dumped", "bytecode", 721, Path(r"C:\r"), None) + keys = self.read_keys(target) + assert keys["bytecode_file"] == "patterns.mlirbc" + assert keys["text_file"] == "" + + def test_first_six_keys_keep_their_documented_order(self, tmp_path: Path) -> None: + target = tmp_path / "metadata.txt" + write_metadata( + target, + make_adapter(0, "Vendor Device", hardware=True), + "dumped", + "text", + 720, + Path(r"C:\r"), + (1, 1), + ) + order = [line.split("=", 1)[0] for line in target.read_text(encoding="utf-8").splitlines()] + assert order[:6] == [ + "adapter", + "driver_version", + "status", + "received_encoding", + "text_file", + "bytecode_file", + ] + + def test_written_with_newline_endings(self, tmp_path: Path) -> None: + target = tmp_path / "metadata.txt" + write_metadata( + target, + make_adapter(0, "Vendor Device", hardware=True), + "dumped", + "text", + 720, + Path(r"C:\r"), + (1, 1), + ) + assert b"\r\n" not in target.read_bytes() + + +# --------------------------------------------------------------------------- +# Cases ported from the original standalone suite (DXML/python/test_cgc.py). +# The MLIR fixtures below are shaped like real driver output rather than the +# minimal sources built above: bodies nest a `rewrite %x { ... }` block, a +# kernel carries its jitFunction *inside* a subgraph_rewrite_desc (so it also +# matches the fusion marker), and one pattern puts a `depends` attribute +# dictionary between the name and the body. +# --------------------------------------------------------------------------- + +KERNEL = ( + b"cgc_pattern.pattern @nvidia.gemm_out : benefit(35) {\n" + b" rewrite %gemm {\n" + b" %d = attribute = #cgc_subgraph_pattern.subgraph_rewrite_desc<\n" + b' jitFunction = "nvidia_gemm_output_major">\n' + b" }\n}\n" +) +FUSION = ( + b"cgc_pattern.pattern @nvidia.expr_native_fusion : benefit(30) {\n" + b' %c = operation "cgc_subgraph_pattern.pattern_cluster"\n' + b" rewrite %c {\n" + b" %d = attribute = #cgc_subgraph_pattern.subgraph_rewrite_desc<\n" + b" cluster_fusion = true>\n" + b" }\n}\n" +) +HINT = ( + b"cgc_pattern.pattern @nvidia.conv_prologue_hint : benefit(50) {\n" + b" rewrite %conv {\n" + b' apply_native_rewrite "cgc_add_pattern_cluster_hint" (%h, %conv)\n' + b" }\n}\n" +) +MARK = ( + b"cgc_pattern.pattern @nvidia.slice_prologue_mark : benefit(0) {\n" + b" rewrite %slice {\n" + b' apply_native_rewrite "cgc_mark_pattern_cluster_op" (%slice, %cat)\n' + b" }\n}\n" +) +REWRITE = ( + b"cgc_pattern.pattern @nvidia.fold_transpose : benefit(1) {\n" + b" rewrite %t {\n" + b" replace %t with %x\n" + b" }\n}\n" +) +WITH_ATTRS = ( + b"cgc_pattern.pattern @nvidia.gemm_epilogue_cluster\n" + b" : benefit(38) attributes {\n" + b" depends = [@nvidia.gemm_expression_epilogue_cluster_hint]\n" + b" } {\n" + b" rewrite %gemm {\n" + b" %d = attribute = #cgc_subgraph_pattern.subgraph_rewrite_desc<\n" + b' jitFunction = "nvidia_gemm_epilogue">\n' + b" }\n}\n" +) +BRANCHING = ( + b"cgc_pattern.pattern @nvidia.rank_agnostic : benefit(36) {\n" + b" any_of {\n" + b" all_of { %a = operand }\n" + b" all_of { %b = operand }\n" + b" all_of { %c = operand }\n" + b" }\n" + b" rewrite %r {\n" + b' apply_native_rewrite "cgc_mark_pattern_cluster_op" (%r, %cat)\n' + b" }\n}\n" +) + +ALL_FIXTURES = KERNEL + FUSION + HINT + MARK + REWRITE + WITH_ATTRS + BRANCHING + + +class TestPatternKindsOnRealisticSources: + """The five kinds, read off bodies shaped like real driver output.""" + + @pytest.mark.parametrize( + ("source", "kind"), + [ + (KERNEL, "kernel"), + (FUSION, "fusion"), + (HINT, "hint"), + (MARK, "mark"), + (REWRITE, "rewrite"), + ], + ) + def test_each_kind_is_recognised(self, source: bytes, kind: str) -> None: + assert split_patterns(source)[0]["kind"] == kind + + def test_a_kernel_is_not_mistaken_for_a_fusion(self) -> None: + # Both carry a subgraph_rewrite_desc; only a kernel names a jitFunction, + # so marker precedence -- not mere presence -- decides the kind. + assert b"subgraph_rewrite_desc" in KERNEL + assert split_patterns(KERNEL)[0]["kind"] == "kernel" + assert split_patterns(FUSION)[0]["kind"] == "fusion" + + def test_name_and_benefit_are_parsed(self) -> None: + record = split_patterns(KERNEL)[0] + assert record["name"] == "nvidia.gemm_out" + assert record["benefit"] == 35 + + def test_a_depends_list_before_the_body_is_not_the_body(self) -> None: + # Taking the first { after the name reads the attribute dictionary as the + # body, finds no jitFunction in it, and files a kernel under rewrite. + record = split_patterns(WITH_ATTRS)[0] + assert record["kind"] == "kernel" + assert record["benefit"] == 38 + + def test_rules_follow_the_any_of_expansion(self) -> None: + assert split_patterns(BRANCHING)[0]["rules"] == 3 + assert split_patterns(MARK)[0]["rules"] == 1 + + def test_two_any_of_blocks_in_one_pattern_multiply(self) -> None: + block = b" any_of {\n all_of { a }\n all_of { b }\n }\n" + src = b"cgc_pattern.pattern @a : benefit(0) {\n" + block + block + b"}\n" + assert count_patterns(src) == (1, 4) + + def test_empty_input(self) -> None: + assert count_patterns(b"") == (0, 0) + + +class TestPatternGroupingAndListing: + """Bucketing by kind and the -v listing built from it.""" + + def test_groups_run_in_the_fixed_kind_order(self) -> None: + groups = pattern_groups(split_patterns(ALL_FIXTURES)) + assert [g["kind"] for g in groups] == ["kernel", "fusion", "hint", "mark", "rewrite"] + + def test_group_carries_its_pattern_and_rule_totals(self) -> None: + groups = pattern_groups(split_patterns(MARK + BRANCHING)) + assert (groups[0]["patterns"], groups[0]["rules"]) == (2, 4) + + def test_patterns_are_listed_by_descending_benefit(self) -> None: + groups = pattern_groups(split_patterns(KERNEL + WITH_ATTRS)) + assert [r["benefit"] for r in groups[0]["members"]] == [38, 35] + + def test_listing_names_every_pattern_exactly_once(self) -> None: + records = split_patterns(ALL_FIXTURES) + text = "\n".join(format_pattern_listing(records)) + for record in records: + assert text.count(record["name"].split(".")[-1]) == 1 + + def test_listing_marks_rule_counts_only_where_they_differ(self) -> None: + text = "\n".join(format_pattern_listing(split_patterns(BRANCHING + MARK))) + assert "3 rules" in text + assert "1 rules" not in text + + +class TestAdapterTable: + """The rendered listing: one row per adapter.""" + + @staticmethod + def render(**overrides: object) -> str: + row = make_adapter(0, "Some GPU", hardware=True) + row.update(overrides) + out = io.StringIO() + print_adapters([row], out) + return out.getvalue() + + @pytest.mark.parametrize(("mlir", "mark"), [(True, "yes"), (False, "no"), (None, "?")]) + def test_support_marks(self, mlir: bool | None, mark: str) -> None: + assert self.render(mlir=mlir).rstrip().endswith(mark) + + def test_unknown_support_prints_a_question_mark(self) -> None: + assert re.search(r"\?\s*$", self.render(mlir=None).splitlines()[-1]) + + def test_header_names_every_column(self) -> None: + header = self.render().splitlines()[0] + assert header.split() == ["IDX", "ADAPTER", "DRIVER", "TYPE", "ATTRIBUTES", "MLIR"] + + def test_a_row_is_written_per_adapter(self) -> None: + rows = [make_adapter(i, f"GPU {i}", hardware=True) for i in range(3)] + out = io.StringIO() + print_adapters(rows, out) + # header, rule, then one line per adapter + assert len(out.getvalue().splitlines()) == 2 + len(rows) + + +class TestAdapterTypeAndAttributes: + """The TYPE and ATTRIBUTES columns.""" + + @staticmethod + def adapter(**overrides: object) -> dict: + base = { + "is_hardware": True, + "is_integrated": True, + "generic_ml": True, + "core_compute": False, + "d3d12_graphics": True, + } + base.update(overrides) + return base + + def test_integrated_hardware(self) -> None: + assert adapter_type(self.adapter()) == "integrated" + + def test_discrete_hardware(self) -> None: + assert adapter_type(self.adapter(is_integrated=False)) == "hardware" + + def test_software_wins_over_integrated(self) -> None: + assert adapter_type(self.adapter(is_hardware=False)) == "software" + + def test_only_advertised_attributes_are_listed(self) -> None: + assert adapter_attrs(self.adapter()) == "ML GFX" + + def test_no_attributes_renders_a_dash(self) -> None: + bare = self.adapter(generic_ml=False, core_compute=False, d3d12_graphics=False) + assert adapter_attrs(bare) == "-" + + +class TestEncodingSniff: + """Telling an MLIR bytecode payload from text.""" + + def test_bytecode_magic_is_recognised(self) -> None: + assert is_mlir_bytecode(b"ML\xefR and then some") + + def test_text_is_not_bytecode(self) -> None: + assert not is_mlir_bytecode(b"cgc_pattern.pattern @a") + + def test_empty_payload_is_not_bytecode(self) -> None: + assert not is_mlir_bytecode(b"") + + def test_a_truncated_magic_is_not_bytecode(self) -> None: + assert not is_mlir_bytecode(b"ML\xef") + + +class TestClearDump: + """--overwrite replaces a dump rather than merging into it.""" + + def test_every_dump_file_is_removed(self, tmp_path: Path) -> None: + for name in DUMP_FILES: + (tmp_path / name).write_text("stale", encoding="utf-8") + (tmp_path / "unrelated.txt").write_text("keep me", encoding="utf-8") + clear_dump(tmp_path) + assert sorted(p.name for p in tmp_path.iterdir()) == ["unrelated.txt"] + + def test_a_stale_bytecode_file_cannot_survive_a_text_dump(self, tmp_path: Path) -> None: + # Without this the directory would describe two different dumps at once. + (tmp_path / "patterns.mlirbc").write_bytes(b"old bytecode") + clear_dump(tmp_path) + assert not (tmp_path / "patterns.mlirbc").exists() + + def test_a_missing_directory_is_harmless(self, tmp_path: Path) -> None: + clear_dump(tmp_path / "no-such-dump-dir") + + +class TestDumpLayout: + """Where a dump is filed.""" + + def test_slug_and_version_directories(self) -> None: + adapter = make_adapter(0, "NVIDIA GeForce RTX 5090 D", hardware=True) + adapter["driver_version"] = "32.0.16.3004" + assert dump_dir(adapter) == Path("patterns") / "nvidia-geforce-rtx-5090-d" / "32.0.16.3004" + + def test_an_unknown_driver_version_keeps_the_layout(self) -> None: + # An empty segment would be joined away, and the dump would lose the level + # that keeps one driver version's capture apart from another's. + adapter = make_adapter(0, "Some GPU", hardware=True) + adapter["driver_version"] = "" + assert dump_dir(adapter) == Path("patterns") / "some-gpu" / "unknown-version" + + def test_the_root_is_always_patterns(self) -> None: + adapter = make_adapter(0, "Intel(R) Graphics", hardware=True) + adapter["driver_version"] = "32.0.101.6127" + assert dump_dir(adapter).parts[0] == "patterns" + + def test_two_vendors_never_collide(self) -> None: + first = make_adapter(0, "NVIDIA GeForce RTX 5090 D", hardware=True) + second = make_adapter(1, "Intel(R) Graphics", hardware=True) + assert dump_dir(first) != dump_dir(second) + + +class TestRedistLine: + """The redist is reported on every run that resolves one.""" + + def test_names_the_directory_and_the_abi(self) -> None: + line = redist_line(Path(r"C:\dxml\build\bin\D3D12"), 720) + assert line.startswith("redist: ") + assert "SDK 720" in line + assert "D3D12" in line + + def test_says_none_when_there_is_no_redist(self) -> None: + line = redist_line(None, None) + assert line.startswith("redist: none") + assert "-v" in line + + +class TestRedistFailureReport: + """Silence about a rejected candidate costs a day.""" + + def test_names_every_candidate_and_refuses_system32(self) -> None: + _, _, tried = resolve_redist(r"C:\definitely-not-here") + text = redist_failure(tried) + assert r"C:\definitely-not-here" in text + assert "missing" in text + assert "System32" in text + assert "does not fall back" in text + + def test_states_the_minimum_sdk(self) -> None: + _, _, tried = resolve_redist(r"C:\definitely-not-here") + assert str(MIN_SDK_VERSION) in redist_failure(tried) + + def test_offers_a_way_out(self) -> None: + _, _, tried = resolve_redist(r"C:\definitely-not-here") + text = redist_failure(tried) + assert "--d3d12-dir" in text + assert REDIST_ENV_VARS[0] in text + + +class TestExchangeStructs: + """The wire layout the driver reads by size.""" + + def test_721_exchange_struct_is_48_bytes(self) -> None: + # 4 bytes of padding sit between Type and IRVersion. A layout that drifts + # is silently wrong rather than loudly wrong. + assert ctypes.sizeof(FeatureDataMLIRExchange721) == 48 + + def test_version_number_packing(self) -> None: + assert format_driver_version(make_version_number(0, 7, 0, 0)) == "0.7.0.0" + + +class TestMetadataConsumerCompatibility: + """Downstream readers parse this file; the shape is a contract.""" + + def test_driver_version_is_findable_by_regex(self, tmp_path: Path) -> None: + # A consumer finds the driver version with exactly this regex. If it stops + # matching, that consumer silently reports a blank driver version. + target = tmp_path / "metadata.txt" + adapter = make_adapter(0, "NVIDIA GeForce RTX 5090 D", hardware=True) + adapter["driver_version"] = "32.0.16.3004" + write_metadata(target, adapter, "dumped", "text", 720, Path(r"C:\r"), (51, 63)) + match = re.search(r"driver_version=(\S+)", target.read_text(encoding="utf-8")) + assert match is not None + assert match.group(1) == "32.0.16.3004" + + def test_every_key_is_present_even_when_empty(self, tmp_path: Path) -> None: + target = tmp_path / "metadata.txt" + write_metadata( + target, + make_adapter(0, "Vendor Device", hardware=True), + "unsupported", + "none", + None, + None, + None, + ) + keys = [line.split("=", 1)[0] for line in target.read_text(encoding="utf-8").splitlines()] + assert len(keys) == len(set(keys)) == 13 + + +# --------------------------------------------------------------------- hardware + +_REDIST, _SDK, _ = resolve_redist() +_NEEDS_REDIST = pytest.mark.skipif( + _REDIST is None, + reason="no Agility SDK redist on this machine (see docs/commands/cgc.md)", +) + + +@_NEEDS_REDIST +class TestCapabilityOnRealHardware: + """Every adapter creates a device; only some drivers answer the exchange.""" + + def test_every_adapter_answers_or_declines(self) -> None: + try: + adapters = list_adapters() + except click.ClickException as exc: # DXCore unavailable + pytest.skip(f"DXCore unavailable: {exc}") + else: + if not adapters: + close_adapters(adapters) + pytest.skip("no D3D12 adapters") + try: + for adapter in adapters: + try: + release(probe_mlir_support(adapter, _REDIST, _SDK or 0)) + except RedistUnusable as exc: # an environment limit, not a code defect + pytest.skip( + f"D3D12 cannot use the redist found here: {exc.format_message()}" + ) + assert adapter["mlir"] in (True, False) + finally: + close_adapters(adapters) + + +# --------------------------------------------------------------------------- +# patterns.json: splitting a dump into one record per pattern. +# +# These fixtures use the syntax real driver dumps actually contain, which differs +# from the minimal sources above in three ways that matter here: the kernel is +# written as #cgc.string<"Name">, source files are separated by "// from" markers, +# and prose comments can also begin with "// from". +# --------------------------------------------------------------------------- + +REAL_KERNEL = ( + b"cgc_pattern.pattern @nvidia.gemm_cluster : benefit(38) {\n" + b" rewrite %gemm {\n" + b" %d = attribute = #cgc_subgraph_pattern.subgraph_rewrite_desc<\n" + b' foreign_config = {jitFunction = #cgc.string<"GemmCluster">}>\n' + b" }\n}\n" +) + + +def merged_dump(*sections: tuple[str, bytes], trailer: bytes = b"") -> bytes: + """Assemble a merged dump the way the driver's does: a manifest, then sections. + + Args: + *sections: ``(source file, pattern sources)`` pairs, in order. + trailer: Bytes appended after the module, such as a NUL terminator. + + Returns: + The dump. + """ + out = [b"// Merged pattern set - do not edit.\n", b"module {\n"] + for name, body in sections: + out.append(b" // from " + name.encode() + b"\n") + out.append(body) + out.append(b"}\n") + return b"".join(out) + trailer + + +class TestSplitPatterns: + """One record per pattern, with the facts a reader needs beyond the counts.""" + + def test_one_record_per_pattern_in_declaration_order(self) -> None: + records = split_patterns(ALL_FIXTURES) + assert [r["index"] for r in records] == list(range(1, len(records) + 1)) + + def test_each_slice_is_exactly_its_own_pattern(self) -> None: + for record in split_patterns(ALL_FIXTURES): + text = record["mlir"] + assert text.startswith(f"cgc_pattern.pattern @{record['name']}") + assert text.rstrip().endswith("}") + assert text.count("cgc_pattern.pattern @") == 1 + assert text.count("{") == text.count("}") + + def test_short_name_drops_the_vendor_prefix(self) -> None: + assert split_patterns(KERNEL)[0]["short_name"] == "gemm_out" + + def test_line_count_matches_the_slice(self) -> None: + record = split_patterns(REAL_KERNEL)[0] + assert record["lines"] == record["mlir"].count("\n") + 1 + + def test_kernel_name_is_read_from_the_real_syntax(self) -> None: + assert split_patterns(REAL_KERNEL)[0]["kernel"] == "GemmCluster" + + def test_kernel_name_accepts_the_plain_string_form(self) -> None: + assert split_patterns(KERNEL)[0]["kernel"] == "nvidia_gemm_output_major" + + def test_non_kernels_declare_no_kernel(self) -> None: + for source in (FUSION, HINT, MARK, REWRITE): + assert split_patterns(source)[0]["kernel"] is None + + def test_a_kernel_named_only_in_a_comment_is_not_a_kernel(self) -> None: + src = MARK.replace( + b" rewrite %slice {", b' // jitFunction = #cgc.string<"Ghost">\n rewrite %slice {' + ) + assert split_patterns(src)[0]["kernel"] is None + + def test_patterns_are_attributed_to_their_source_file(self) -> None: + dump = merged_dump(("a.pdll", KERNEL + MARK), ("b.mlir", HINT)) + assert [r["source"] for r in split_patterns(dump)] == ["a.pdll", "a.pdll", "b.mlir"] + + def test_a_prose_from_comment_does_not_open_a_section(self) -> None: + prose = b" // from outside the root's own chain. Both hints create clusters\n" + dump = merged_dump(("a.pdll", KERNEL + prose + MARK)) + assert {r["source"] for r in split_patterns(dump)} == {"a.pdll"} + + def test_a_pattern_before_any_section_is_unattributed(self) -> None: + assert split_patterns(KERNEL)[0]["source"] is None + + def test_the_trailing_nul_is_dropped_before_splitting(self) -> None: + # The driver reports its payload size including the C string terminator. + dump = normalize_text_payload(merged_dump(("a.pdll", KERNEL), trailer=b"\x00")) + assert not dump.endswith(b"\x00") + records = split_patterns(dump) + assert len(records) == 1 + assert "\x00" not in records[0]["mlir"] + + def test_a_declaration_without_a_body_does_not_adopt_the_next_one(self) -> None: + # The body search stops at the next declaration, so a body-less one is skipped + # rather than swallowing the pattern that follows it. + source = b"cgc_pattern.pattern @p.nobody\n" + KERNEL + records = split_patterns(source) + assert [r["name"] for r in records] == [r["name"] for r in split_patterns(KERNEL)] + assert records[0]["mlir"].count("cgc_pattern.pattern @") == 1 + + def test_a_comment_marker_inside_a_string_is_not_a_comment(self) -> None: + # Blanking from "//" to end of line would take the closing brace with it, and + # the pattern would swallow the next one. + source = ( + b"cgc_pattern.pattern @p.one : benefit(1) {\n" + b' jitFunction = "http://example/kernel" }\n' + ) + KERNEL + records = split_patterns(source) + assert [r["short_name"] for r in records] == [ + "one", + split_patterns(KERNEL)[0]["short_name"], + ] + assert records[0]["lines"] == 2 + + def test_comments_survive_in_the_displayed_text(self) -> None: + src = KERNEL.replace(b" rewrite %gemm {", b" // keep me\n rewrite %gemm {") + assert "// keep me" in split_patterns(src)[0]["mlir"] + + def test_a_brace_inside_a_comment_cannot_unbalance_the_slice(self) -> None: + src = KERNEL.replace(b" rewrite %gemm {", b" // stray } brace\n rewrite %gemm {") + record = split_patterns(src + MARK)[0] + assert record["name"] == "nvidia.gemm_out" + assert record["mlir"].count("cgc_pattern.pattern @") == 1 + assert record["kind"] == "kernel" + + +class TestParserEquivalence: + """Blanking comments must count exactly what deleting them did.""" + + @staticmethod + def deleting_parser(data: bytes) -> tuple[int, int]: + """The parser's previous behaviour: delete comments, then count.""" + return count_patterns(COMMENT_RE.sub(b"", data)) + + @pytest.mark.parametrize( + "source", + [ALL_FIXTURES, KERNEL + MARK, BRANCHING, WITH_ATTRS, REAL_KERNEL], + ids=["all", "kernel+mark", "branching", "depends", "real-kernel"], + ) + def test_counts_are_unchanged(self, source: bytes) -> None: + assert count_patterns(source) == self.deleting_parser(source) + + def test_counts_are_unchanged_with_comments_everywhere(self) -> None: + noisy = b"\n".join(line + b" // trailing note" for line in ALL_FIXTURES.split(b"\n")) + assert count_patterns(noisy) == self.deleting_parser(noisy) + + +class TestPatternsDocument: + """The patterns.json document written beside a dump.""" + + @staticmethod + def document(data: bytes) -> dict: + adapter = make_adapter(0, "NVIDIA GeForce RTX 5090 D", hardware=True) + return patterns_document(data, adapter, 720, Path(r"C:\redist")) + + def test_carries_its_format_tag(self) -> None: + assert self.document(ALL_FIXTURES)["format"] == PATTERNS_JSON_FORMAT + + def test_meta_records_provenance(self) -> None: + meta = self.document(ALL_FIXTURES)["meta"] + assert meta["adapter"] == "NVIDIA GeForce RTX 5090 D" + assert meta["sdk_version"] == 720 + assert meta["bytes"] == len(ALL_FIXTURES) + + def test_summary_agrees_with_the_patterns_it_carries(self) -> None: + doc = self.document(ALL_FIXTURES) + pats, summary = doc["patterns"], doc["summary"] + assert summary["patterns"] == len(pats) + assert summary["rules"] == sum(p["rules"] for p in pats) + assert summary["kernel_patterns"] == sum(1 for p in pats if p["kernel"]) + assert summary["top_benefit"] == max(p["benefit"] or 0 for p in pats) + + def test_summary_agrees_with_count_patterns(self) -> None: + summary = self.document(ALL_FIXTURES)["summary"] + assert (summary["patterns"], summary["rules"]) == count_patterns(ALL_FIXTURES) + + def test_every_pattern_is_in_exactly_one_source_group(self) -> None: + doc = self.document(merged_dump(("a.pdll", KERNEL + MARK), ("b.mlir", HINT))) + grouped = [i for s in doc["sources"] for i in s["patterns"]] + assert sorted(grouped) == [p["index"] for p in doc["patterns"]] + assert [s["name"] for s in doc["sources"]] == ["a.pdll", "b.mlir"] + + def test_kind_totals_sum_to_the_pattern_count(self) -> None: + doc = self.document(ALL_FIXTURES) + assert sum(doc["kinds"].values()) == doc["summary"]["patterns"] + + def test_distinct_kernels_counts_each_name_once(self) -> None: + doc = self.document(REAL_KERNEL + REAL_KERNEL.replace(b"gemm_cluster", b"gemm_cluster_2")) + assert doc["summary"]["kernel_patterns"] == 2 + assert doc["summary"]["distinct_kernels"] == 1 + + +class TestWritePatternsJson: + """The file on disk.""" + + def test_round_trips_through_json(self, tmp_path: Path) -> None: + adapter = make_adapter(0, "Vendor Device", hardware=True) + doc = patterns_document(ALL_FIXTURES, adapter, 720, Path(r"C:\r")) + target = tmp_path / "patterns.json" + write_patterns_json(target, doc) + assert json.loads(target.read_text(encoding="utf-8")) == doc + + def test_written_with_newline_endings(self, tmp_path: Path) -> None: + adapter = make_adapter(0, "Vendor Device", hardware=True) + target = tmp_path / "patterns.json" + write_patterns_json(target, patterns_document(KERNEL, adapter, 720, None)) + assert b"\r\n" not in target.read_bytes() + + def test_is_part_of_a_dump_so_overwrite_clears_it(self, tmp_path: Path) -> None: + assert "patterns.json" in DUMP_FILES + (tmp_path / "patterns.json").write_text("{}", encoding="utf-8") + clear_dump(tmp_path) + assert not (tmp_path / "patterns.json").exists() + + +class TestAtlasPage: + """--open shows a dump in the viewer that ships with the docs.""" + + @staticmethod + def seed_of(path: Path) -> str: + text = path.read_text(encoding="utf-8") + start = text.index('id="seed">') + len('id="seed">') + return text[start : text.index("", start)] + + @pytest.fixture + def opened(self, monkeypatch: pytest.MonkeyPatch) -> list[str]: + seen: list[str] = [] + monkeypatch.setattr("webbrowser.open", lambda uri: bool(seen.append(uri)) or True) + return seen + + def test_the_page_ships_with_the_docs(self) -> None: + assert atlas_page().name == "cgc-pattern-atlas.html" + + def test_without_a_dump_the_page_itself_is_opened(self, opened: list[str]) -> None: + open_atlas(None) + assert opened == [atlas_page().as_uri()] + + def test_a_dump_is_embedded_so_the_page_opens_on_it( + self, tmp_path: Path, opened: list[str] + ) -> None: + adapter = make_adapter(0, "NVIDIA GeForce RTX 5090 D", hardware=True) + document = tmp_path / "patterns.json" + write_patterns_json(document, patterns_document(ALL_FIXTURES, adapter, 720, None)) + + open_atlas(document) + + copy = path_of(opened[0]) + seeded = json.loads(self.seed_of(copy)) + assert seeded["patterns"] == json.loads(document.read_text(encoding="utf-8"))["patterns"] + assert seeded["meta"]["label"] == str(document) + + def test_open_with_a_file_shows_it_without_touching_any_driver( + self, tmp_path: Path, opened: list[str] + ) -> None: + # Needs neither a redist nor a GPU: run_patterns returns before DXCore. + adapter = make_adapter(0, "NVIDIA GeForce RTX 5090 D", hardware=True) + document = tmp_path / "patterns.json" + write_patterns_json(document, patterns_document(ALL_FIXTURES, adapter, 720, None)) + + run_patterns(Options(open_atlas=str(document))) + + seeded = json.loads(self.seed_of(path_of(opened[0]))) + assert seeded["meta"]["label"] == str(document) + + def test_open_with_a_missing_file_is_a_usage_error(self, tmp_path: Path) -> None: + with pytest.raises(click.UsageError): + run_patterns(Options(open_atlas=str(tmp_path / "absent.json"))) + + def test_open_with_a_file_and_dump_together_is_a_usage_error(self, tmp_path: Path) -> None: + # Two dumps to show; preferring either silently throws away what was asked for. + with pytest.raises(click.UsageError): + run_patterns(Options(adapter="nvidia", dump=True, open_atlas=str(tmp_path / "p.json"))) + + @pytest.mark.parametrize( + ("name", "content"), + [ + ("notes.txt", "not json at all"), + ("other.json", '{"format": "something-else/1"}'), + ("bare.json", "[]"), + ], + ) + def test_open_refuses_a_file_that_is_not_a_dump( + self, tmp_path: Path, name: str, content: str + ) -> None: + # A mistyped path that happens to exist is a bad argument, not a traceback. + wrong = tmp_path / name + wrong.write_text(content, encoding="utf-8") + with pytest.raises(click.UsageError): + run_patterns(Options(open_atlas=str(wrong))) + + def test_a_pattern_cannot_close_the_script_element( + self, tmp_path: Path, opened: list[str] + ) -> None: + # Every "<" is escaped, so a pattern whose own text holds -- which + # would end the element early and break the page -- arrives intact. + document = tmp_path / "patterns.json" + hostile = { + "format": PATTERNS_JSON_FORMAT, + "meta": {}, + "patterns": [{"mlir": ""}], + } + document.write_text(json.dumps(hostile), encoding="utf-8") + + open_atlas(document) + + raw = self.seed_of(path_of(opened[0])) + assert "<" not in raw + assert json.loads(raw)["patterns"] == hostile["patterns"]