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/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/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 a8763a046..112d41589 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" ] @@ -176,12 +176,16 @@ 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" = [ "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/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"] 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")