diff --git a/docs/commands/compile.md b/docs/commands/compile.md index c57fa63f9..37a0cd5a9 100644 --- a/docs/commands/compile.md +++ b/docs/commands/compile.md @@ -1,6 +1,6 @@ # winml compile -> Compile an ONNX model to an EP-specific format for fast runtime loading. +> Compile ONNX to an EP-specific format, or explicitly compile CGIR to device-targeted CG Output IR. ## When to use this @@ -18,14 +18,15 @@ $ winml compile [options] | Flag | Short | Type | Default | Description | |---|---|---|---|---| -| `--model` | `-m` | path | *(required unless `--list`)* | Input ONNX model file. | +| `--model` | `-m` | path | *(required unless `--list`)* | Input ONNX model, or complete CGIR in explicit IR mode. | | `--output` | `-o` | path | — | Output file path (e.g., `model_compiled.onnx`). Takes precedence over `--output-dir`. | | `--output-dir` | | path | same dir as input | Directory to write compiled output artifacts. | | `--device` | `-d` | choice | `auto` | Target device: `auto`, `npu`, `gpu`, or `cpu`. | +| `--device-luid` | | string | — | Explicit CGC IR mode only: GPU adapter LUID from `winml sys`, in `0xHHHHHHHH_0xLLLLLLLL` form. | | `--ep` | | `TEXT` | — | Force a specific execution provider, overriding device-to-provider mapping. Accepts full names (e.g., `QNNExecutionProvider`) or aliases (`qnn`, `dml`, `openvino`, `vitisai`, `migraphx`, `cpu`, `nvtensorrtrtx`). | | `--ep-options` | | `KEY=VALUE` | — | Execution-provider option for the compilation session. Repeat for multiple options; later duplicate keys win. Explicit CLI values override matching `compile.provider_options` keys from `--config`. | -| `--validate` / `--no-validate` | | flag | `--validate` | Run a post-compilation validation pass on the target hardware. Enabled by default; pass `--no-validate` to skip when the target hardware or driver is unavailable. | -| `--compiler` | | choice | `ort` | Compiler backend: `ort` (ONNX Runtime) or `qairt` (Qualcomm AI Runtime Tools). | +| `--validate` / `--no-validate` | | flag | `--validate` | EPContext: post-compilation inference validation. CG Output IR: reload, pipeline build and I/O schema validation, without inference. | +| `--compiler` | | choice | `ort` | Backends: `ort`, `ort_session`, `qairt`, `cgc`. `cgc` is inferred for a single `.mlir` or `.mlirbc` input. | | `--qnn-sdk-root` | | path | `None` | Path to the QNN SDK root directory. | | `--embed/--no-embed` | | flag | `false` | Embed the EP context blob inside the ONNX file instead of writing a separate `.bin` file. | | `--list` | | flag | `false` | List available compiler backends for the selected device and exit without compiling. | @@ -48,6 +49,61 @@ If a provider option names a compiler input file, list its key in that option path and fingerprints its contents for EPContext cache identity; other provider-option strings are always passed through unchanged. +## CGIR to CG Output IR + +This explicit mode uses the Windows ML Runtime API directly, without an ONNX +Runtime EP. The input and output selectors make the representation unambiguous: + +```powershell +winml compile ` + -m input.mlir ` + --compiler cgc ` + --device gpu ` + -o output.mlir +``` + +`--compiler cgc` is optional for `.mlir` and `.mlirbc` input. Specify it when +the CGIR file uses another suffix or when you want to make the backend explicit. + +The command loads CGIR through Windows ML Runtime. Both textual and bytecode +CGIR are accepted; invalid or unsupported input is reported by the Runtime. +Supply a complete executable model with its weights, not a graph-only report +projection. + +Compilation calls the selected GPU target's Runtime +`ModelCompiler.compile_to_file()` with `CompiledModelForm.DEVICE_TARGETED`. +The saved artifact is **MLIR bytecode**, even when named `.mlir`; text dumping +is not provided by this command. `.mlirbc` is also accepted as an output suffix. +No optimization or quantization pipeline runs. + +This first version supports GPU only (the default in explicit IR mode). A single +GPU is selected automatically; multiple GPUs require `--device-luid`. GPU +discovery uses DXCore, not EP discovery or provider installation. + +Default `--validate` reloads the compiled artifact, builds a Runtime pipeline and +checks that its I/O schema matches the input. Unlike the EPContext validation +path, this does not execute inference or certify accuracy. `--no-validate` skips +that reload/build check; genuine Input IR parsing and compilation remain required. +Run `winml perf` separately for performance. + +The command publishes the artifact and any Runtime-produced resource companions +only after successful compilation and requested validation. It preserves an +existing input's `_metadata.json` I/O bindings as output metadata and writes +`_compile.json` with input/artifact hashes, output encoding, GPU LUID +and validation status. Metadata and resource collisions also require +`--overwrite`; the input and its metadata are never replaced. Without `-o`, the +default is `_output.mlir` in `--output-dir` or beside the input. + +`--ep`, `--ep-options`, `--embed`/`--no-embed`, `--qnn-sdk-root`, `--list`, +`--config` and multiple `-m` inputs are rejected in this mode. Use the existing +ONNX-to-EPContext mode for EP compilation. Omitting the new selectors preserves +the original command behavior. + +The preview `windowsml` payload must include Runtime APIs compatible with the +active CLI. If preview wheels override stable project pins, invoke the +environment directly or use `uv run --no-sync` to avoid restoring stable +packages. + ## Examples ```bash diff --git a/docs/getting-started/cgc-onboarding.md b/docs/getting-started/cgc-onboarding.md index c154289f7..4dfba1dc7 100644 --- a/docs/getting-started/cgc-onboarding.md +++ b/docs/getting-started/cgc-onboarding.md @@ -144,6 +144,21 @@ winml export -m .\model.onnx --target cgir -o .\model.mlir The model must already be compatible with the CGC converter. +### CGIR to device-targeted CG Output IR + +To retain a compiled artifact from existing CGIR, use the explicit +[CGIR compile mode](../commands/compile.md#cgir-to-cg-output-ir): + +```powershell +winml compile -m .\model.mlir --compiler cgc --device gpu -o .\model-output.mlir +``` + +For `.mlir` and `.mlirbc` input, `--compiler cgc` is optional. + +The artifact is bytecode, not textual MLIR. This is a separate compilation stage, +without an EP, optimization or quantization. Default validation reloads/builds +the artifact and checks I/O; it does not run inference or performance tests. + ## Tutorials and samples - [ResNet-50 - PyTorch to CGIR (experimental)](../samples/resnet50-cgir.md) walks through diff --git a/src/winml/modelkit/commands/compile.py b/src/winml/modelkit/commands/compile.py index defc9a8a2..85129df5d 100644 --- a/src/winml/modelkit/commands/compile.py +++ b/src/winml/modelkit/commands/compile.py @@ -40,6 +40,8 @@ if TYPE_CHECKING: + from typing import Literal, TypeGuard + from ..utils.constants import CompilerName, EPName, EPNameOrAlias from ..utils.logging import configure_logging from ._ep_arg import EpAtSourceParamType @@ -62,11 +64,16 @@ def _resolve_compile_provider(resolved_device: str, ep: EPNameOrAlias | None) -> console = Console() +def _is_epcontext_compiler(value: str) -> TypeGuard[CompilerName]: + return value in COMPILER_NAMES + + @click.command() @cli_utils.model_path_option( required=False, multiple=True, - help_text="Input ONNX model file. Repeat -m to compile multiple models with a shared " + help_text="Input model file (ONNX by default; CGIR input is also supported). " + "Repeat -m to compile multiple ONNX models with a shared " "EP context (weight sharing). Required unless --list.", ) @cli_utils.output_option("Output file path (e.g., model_compiled.onnx)") @@ -84,6 +91,9 @@ def _resolve_compile_provider(resolved_device: str, ep: EPNameOrAlias | None) -> default=None, help="Target device (default: deduced from --ep, or 'npu' if neither given)", ) +@cli_utils.device_luid_option( + optional_message="Only supported for CGIR to CG Output IR compilation." +) @click.option( "--ep", type=EpAtSourceParamType(), @@ -98,14 +108,16 @@ def _resolve_compile_provider(resolved_device: str, ep: EPNameOrAlias | None) -> @click.option( "--validate/--no-validate", default=True, - help="Validate compiled model (default: enabled)", + help="Validate compiled model. EPContext: inference; CG Output IR: reload, pipeline " + "build and I/O schema (no inference).", ) @click.option( "--compiler", - type=click.Choice(list(COMPILER_NAMES)), + type=click.Choice([*COMPILER_NAMES, "cgc"]), default="ort", help="Compiler backend (default: ort). 'ort_session' compiles via " - "ort.InferenceSession (ep.context_enable) — required for shared-context multi-model.", + "ort.InferenceSession (ep.context_enable) — required for shared-context multi-model. " + "'cgc' compiles CGIR through Windows ML Runtime and is inferred for .mlir/.mlirbc input.", ) @click.option( "--qnn-sdk-root", @@ -137,18 +149,19 @@ def compile( output_dir: Path | None, overwrite: bool, device: str | None, + device_luid: str | None, ep: tuple[str, str | None] | None, ep_options: tuple[str, ...], validate: bool, verbose: int, quiet: bool, - compiler: CompilerName, + compiler: CompilerName | Literal["cgc"], qnn_sdk_root: Path | None, embed: bool, list_compilers_flag: bool, config_file: Path | None, ) -> None: - r"""Compile ONNX model to EP-specific format. + r"""Compile ONNX to EPContext, or CGIR to CG Output IR. This command compiles an ONNX model to an EP-specific format (e.g., QNN EPContext). @@ -166,10 +179,64 @@ def compile( # Compile using QAIRT SDK winml compile -m model.onnx --compiler qairt --qnn-sdk-root /path/to/sdk + + # Compile CGIR to device-targeted bytecode CG Output IR + winml compile -m input.mlir --compiler cgc --device gpu -o output.mlir """ # Merge top-level -v/-q with subcommand-level flags so either position works. verbose, quiet = cli_utils.resolve_verbosity(ctx, verbose, quiet) + compiler_was_provided = cli_utils.is_cli_provided(ctx, "compiler") + inferred_cgc = ( + not compiler_was_provided + and len(model) == 1 + and model[0].suffix.casefold() in {".mlir", ".mlirbc"} + ) + ir_mode = compiler == "cgc" or inferred_cgc + if ir_mode: + incompatible = { + "--ep": ep is not None, + "--ep-options": bool(ep_options), + "--qnn-sdk-root": qnn_sdk_root is not None, + "--embed/--no-embed": cli_utils.is_cli_provided(ctx, "embed"), + "--list": list_compilers_flag, + "--config": config_file is not None, + } + rejected = [name for name, present in incompatible.items() if present] + if rejected: + raise click.UsageError( + f"CG Output IR compilation does not support: {', '.join(rejected)}." + ) + if len(model) != 1: + raise click.UsageError("CG Output IR compilation requires exactly one --model.") + if device not in {None, "gpu"}: + raise click.UsageError("CG Output IR compilation currently requires --device gpu.") + from ..compiler import compile_cgc_ir + + source = model[0] + destination = output or (output_dir or source.parent) / f"{source.stem}_output.mlir" + configure_logging(verbosity=verbose, quiet=quiet) + try: + result_path = compile_cgc_ir( + source, + destination, + device_luid=device_luid, + validate=validate, + overwrite=overwrite, + ) + except OSError as e: + raise click.ClickException(f"CGC compilation failed: {e}") from e + console.print(f"Output: {result_path}") + console.print("Format: MLIR bytecode; form: DEVICE_TARGETED (CG Output IR).") + console.print( + "Validation: reload, pipeline build and I/O schema passed (no inference)." + if validate + else "Validation: skipped by --no-validate." + ) + return + if device_luid is not None: + raise click.UsageError("--device-luid is only supported for CG Output IR compilation.") + # --ep is parsed by EpAtSourceParamType at click parse time and # arrives as (ep, source) or None — feeds the EPDeviceTarget's # `source` axis (Scenarios A.5/A.6 per 2_coreloop.md §6.2). Unpack @@ -227,6 +294,12 @@ def compile( if verbose == 0 and not quiet and "verbose" in cc: verbose = int(cc["verbose"]) + if not _is_epcontext_compiler(compiler): + raise click.UsageError( + "CGC compilation requires a single .mlir/.mlirbc model or explicit " + "--compiler cgc; it cannot be selected through EPContext configuration." + ) + configure_logging(verbosity=verbose, quiet=quiet) # Resolve EP+device at the CLI boundary (plan §C / Decision B). diff --git a/src/winml/modelkit/compiler/__init__.py b/src/winml/modelkit/compiler/__init__.py index b30124446..e8ab53337 100644 --- a/src/winml/modelkit/compiler/__init__.py +++ b/src/winml/modelkit/compiler/__init__.py @@ -47,6 +47,7 @@ # triggering the heavy imports at runtime. if TYPE_CHECKING: from .compiler import Compiler, compile_multiple_onnx, compile_onnx, list_compilers + from .runtime import compile_cgc_ir from .stages.compile import CompileStage from .stages.optimize import OptimizeStage from .stages.qformat import QFormatConvertStage @@ -54,6 +55,11 @@ def __getattr__(name: str) -> Any: """Lazy-load heavy symbols that pull in session/torch to speed up import.""" + if name == "compile_cgc_ir": + from .runtime import compile_cgc_ir + + globals()[name] = compile_cgc_ir + return compile_cgc_ir if name in {"Compiler", "compile_multiple_onnx", "compile_onnx", "list_compilers"}: from .compiler import ( Compiler, @@ -99,6 +105,7 @@ def __getattr__(name: str) -> Any: "QFormatConvertStage", "WinMLCompileConfig", "clear_transforms", + "compile_cgc_ir", "compile_multiple_onnx", "compile_onnx", "get_transforms_for_ep", diff --git a/src/winml/modelkit/compiler/runtime.py b/src/winml/modelkit/compiler/runtime.py new file mode 100644 index 000000000..80f5883d0 --- /dev/null +++ b/src/winml/modelkit/compiler/runtime.py @@ -0,0 +1,192 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Device-targeted compilation of existing CGC Input IR through Runtime.""" + +from __future__ import annotations + +import hashlib +import json +import logging +import time +from pathlib import Path +from tempfile import TemporaryDirectory +from typing import TYPE_CHECKING + +import click +from filelock import FileLock + +from ..export.cgc.artifacts import cgc_metadata_path +from ..session.runtime_session import WinMLRuntimeSession +from ..sysinfo.dxcore_adapters import enumerate_compute_adapters +from ..utils.cli import guard_output + + +if TYPE_CHECKING: + from ..sysinfo.dxcore_adapters import DXCoreAdapterInfo + +logger = logging.getLogger(__name__) + + +def _file_identity(path: Path) -> dict[str, str | int]: + with path.open("rb") as stream: + digest = hashlib.file_digest(stream, "sha256").hexdigest() + return {"filename": path.name, "size_bytes": path.stat().st_size, "sha256": digest} + + +def _select_gpu(device_luid: str | None) -> DXCoreAdapterInfo: + adapters = [adapter for adapter in enumerate_compute_adapters() if adapter.device_type == "GPU"] + if device_luid is not None: + adapters = [ + adapter for adapter in adapters if adapter.luid.casefold() == device_luid.casefold() + ] + if len(adapters) != 1: + raise click.ClickException( + "CGC compilation requires one unambiguous GPU. Use --device-luid with " + f"a current GPU LUID from 'winml sys'; found {len(adapters)} matching adapters." + ) + return adapters[0] + + +def _same_file(first: Path, second: Path) -> bool: + return first.resolve() == second.resolve() or ( + first.exists() and second.exists() and first.samefile(second) + ) + + +def _publish( + staging: Path, + output: Path, + overwrite: bool, + protected: tuple[Path, ...], + invalidated: tuple[Path, ...] = (), +) -> None: + files = [path for path in staging.iterdir() if path.is_file()] + if any(path.is_dir() for path in staging.iterdir()): + raise click.ClickException("Cannot publish a nested Runtime artifact layout.") + destinations = [output.parent / path.name for path in files] + stale = [path for path in invalidated if path not in destinations and path.exists()] + for destination in [*destinations, *stale]: + if any(_same_file(destination, source) for source in protected): + raise click.ClickException(f"Output would overwrite a source artifact: {destination}") + if destination.is_symlink() or (destination.exists() and not destination.is_file()): + raise click.ClickException(f"Output artifact is not a regular file: {destination}") + guard_output(destination, overwrite) + backup = staging / "backup" + backup.mkdir() + restored: list[tuple[Path, Path]] = [] + published: list[Path] = [] + try: + for index, destination in enumerate(stale): + saved = backup / f"stale-{index}" + destination.replace(saved) + restored.append((saved, destination)) + for index, path in enumerate(files): + destination = output.parent / path.name + if destination.exists(): + saved = backup / str(index) + destination.replace(saved) + restored.append((saved, destination)) + path.replace(destination) + published.append(destination) + except BaseException: + for destination in reversed(published): + destination.unlink() + for saved, destination in reversed(restored): + saved.replace(destination) + raise + + +def compile_cgc_ir( + model_path: Path, + output_path: Path, + *, + device_luid: str | None = None, + validate: bool = True, + overwrite: bool = False, +) -> Path: + """Compile genuine CGC Input IR to a persistent device-targeted bytecode bundle. + + Validation reloads the artifact, builds a pipeline and compares I/O schemas; + it does not run inference, numerical accuracy checks or performance tests. + """ + started = time.perf_counter() + if not model_path.is_file(): + raise click.ClickException(f"Input IR file not found: {model_path}") + if _same_file(model_path, output_path): + raise click.ClickException( + "Input and output must be different files, even with --overwrite." + ) + if output_path.suffix.lower() not in {".mlir", ".mlirbc"}: + raise click.UsageError("CG Output IR output must use .mlir or .mlirbc.") + metadata_path = cgc_metadata_path(model_path) + metadata_identity = _file_identity(metadata_path) if metadata_path.is_file() else None + metadata = None + if metadata_path.is_file(): + try: + metadata = json.loads(metadata_path.read_text(encoding="utf-8")) + except (ValueError, UnicodeError) as e: + raise click.ClickException(f"Invalid CGC metadata '{metadata_path}': {e}") from e + if not isinstance(metadata, dict): + raise click.ClickException(f"CGC metadata must be an object: {metadata_path}") + guard_output(output_path, overwrite) + source_identity = _file_identity(model_path) + adapter_info = _select_gpu(device_luid) + high, low = adapter_info.luid.split("_") + luid = (int(high, 16) << 32) | int(low, 16) + logger.info("Compiling CGC Input IR on %s / %s", adapter_info.name, adapter_info.luid) + output_path.parent.mkdir(parents=True, exist_ok=True) + with ( + FileLock(str(output_path) + ".compile.lock"), + TemporaryDirectory( + prefix=f".{output_path.stem}-compile-", dir=output_path.parent + ) as temporary, + ): + staging = Path(temporary) + artifact = staging / output_path.name + session = WinMLRuntimeSession(model_path, backend="cgc", device_luid=luid) + session.compile_to_file(artifact, validate=validate) + if not artifact.is_file() or not artifact.stat().st_size: + raise click.ClickException("Runtime returned no nonempty Output IR artifact.") + with artifact.open("rb") as stream: + if stream.read(4) != b"ML\xefR": + raise click.ClickException("Runtime did not produce MLIR bytecode Output IR.") + if metadata is not None: + cgc_metadata_path(artifact).write_text( + json.dumps({**metadata, "format": "cg-output-ir"}, indent=2), + encoding="utf-8", + ) + if _file_identity(model_path) != source_identity: + raise click.ClickException("Input IR changed during compilation; output not published.") + current_metadata_identity = ( + _file_identity(metadata_path) if metadata_path.is_file() else None + ) + if current_metadata_identity != metadata_identity: + raise click.ClickException("Input IR metadata changed; output not published.") + receipt = { + "input_format": "cgir", + "target": "cg-output-ir", + "runtime": "winml-runtime", + "compiled_form": "DEVICE_TARGETED", + "output_encoding": "mlir-bytecode", + "input": source_identity, + "gpu": {"name": adapter_info.name, "luid": adapter_info.luid}, + "validation": { + "status": "passed" if validate else "skipped", + "scope": "reload, pipeline build and I/O schema; no inference or accuracy", + }, + "artifacts": [_file_identity(path) for path in staging.iterdir() if path.is_file()], + "elapsed_seconds": time.perf_counter() - started, + } + (staging / f"{output_path.stem}_compile.json").write_text( + json.dumps(receipt, indent=2), encoding="utf-8" + ) + _publish( + staging, + output_path, + overwrite, + (model_path, metadata_path), + invalidated=(cgc_metadata_path(output_path),), + ) + return output_path diff --git a/src/winml/modelkit/session/runtime_session.py b/src/winml/modelkit/session/runtime_session.py index 3f19d8db8..79cd00a55 100644 --- a/src/winml/modelkit/session/runtime_session.py +++ b/src/winml/modelkit/session/runtime_session.py @@ -499,6 +499,15 @@ def _resolve_mlir_target( return _ResolvedRuntimeTarget(target, device_class) adapter_luid = device.adapter_luid + return _resolve_adapter_target(runtime, adapter_luid, device_class) + + +def _resolve_adapter_target( + runtime: Any, + adapter_luid: int | None, + device_class: str = "gpu", +) -> _ResolvedRuntimeTarget: + """Create a Runtime target directly from a DXCore adapter LUID.""" adapter = None adapter_error = None if adapter_luid is not None: @@ -520,6 +529,52 @@ def _resolve_mlir_target( raise +def _compile_runtime_model( + source_model: Any, + resolved_target: _ResolvedRuntimeTarget, + output_path: Path, + *, + symbolic_dimensions: Mapping[str, int] | None = None, + compiled_form: Any | None = None, +) -> None: + """Compile a loaded Runtime model while owning compiler/options lifetimes.""" + with _translate_native_errors("build", device_class=resolved_target.device_class): + compiler = resolved_target.execution_target.model_compiler() + options = None + try: + if symbolic_dimensions or compiled_form is not None: + options = compiler.create_options() + if symbolic_dimensions: + if not options.supports_symbolic_dimensions: + raise click.ClickException( + "The installed Runtime compiler does not support symbolic dimensions." + ) + for name, extent in symbolic_dimensions.items(): + options.symbolic_dimensions[name] = extent + if compiled_form is not None: + if not compiler.supports_form(compiled_form): + raise click.ClickException( + "Runtime does not support DEVICE_TARGETED CG Output IR." + ) + options.form = compiled_form + compiler.compile_to_file(source_model, str(output_path), options=options) + else: + compiler.compile_to_file(source_model, str(output_path)) + finally: + if options is not None: + options.close() + compiler.close() + + +def _normalized_schema_descriptors(schema: Any) -> tuple[list[Any], list[Any]]: + """Return comparable I/O descriptors with Runtime dynamic dims normalized.""" + + def normalize(items: Any) -> list[Any]: + return [(dtype, tuple(_shape_with_dynamic_dims(shape))) for dtype, shape in items] + + return normalize(schema.inputs()), normalize(schema.outputs()) + + def _resolve_onnx_target( runtime: Any, wr: Any, @@ -638,6 +693,7 @@ def __init__( provider_options: Mapping[str, str] | None = None, session_options: Callable[[], Any] | None = None, backend: RuntimeBackend, + device_luid: int | None = None, ) -> None: """Initialize a Runtime session. @@ -653,17 +709,22 @@ def __init__( """ self._lock = threading.RLock() self._model_path = Path(model_path) - self._is_mlir = self._model_path.suffix.lower() == ".mlir" + self._is_mlir = self._model_path.suffix.lower() in {".mlir", ".mlirbc"} self._backend = backend - if self._backend == "cgc" and ep_device is None: - raise ValueError("ep_device is required for CGC Runtime sessions.") + if self._backend == "cgc" and ep_device is None and device_luid is None: + raise ValueError("ep_device or device_luid is required for CGC Runtime sessions.") self._ep_device = ep_device - if self._backend == "cgc": - assert ep_device is not None + self._device_luid = device_luid + if self._backend == "cgc" and ep_device is not None: device = ep_device.device.device_type.lower() ep = None ep_source = None provider_options = None + elif self._backend == "cgc": + device = "gpu" + ep = None + ep_source = None + provider_options = None elif ep_device is not None: device = ep_device.device.device_type.lower() ep = ep_device.ep_short_name @@ -723,8 +784,9 @@ def _ensure_built(self) -> None: def _resolve_target(self, runtime: Any, wr: Any) -> _ResolvedRuntimeTarget: if self._backend == "cgc": - assert self._ep_device is not None - return _resolve_mlir_target(runtime, self._ep_device) + if self._ep_device is not None: + return _resolve_mlir_target(runtime, self._ep_device) + return _resolve_adapter_target(runtime, self._device_luid) return _resolve_onnx_target( runtime, wr, @@ -753,26 +815,12 @@ def _load_onnx_on_cgc( self._compiled_artifacts = TemporaryDirectory(prefix="winml-runtime-cgc-") artifact_path = Path(self._compiled_artifacts.name) / "model.mlir" - with _translate_native_errors("build", device_class="gpu"): - compiler = resolved_target.execution_target.model_compiler() - options = None - try: - if self._symbolic_dimensions: - options = compiler.create_options() - if not options.supports_symbolic_dimensions: - raise click.ClickException( - "The installed Runtime compiler does not support " - "symbolic dimensions." - ) - for name, extent in self._symbolic_dimensions.items(): - options.symbolic_dimensions[name] = extent - compiler.compile_to_file(source_model, str(artifact_path), options=options) - else: - compiler.compile_to_file(source_model, str(artifact_path)) - finally: - if options is not None: - options.close() - compiler.close() + _compile_runtime_model( + source_model, + resolved_target, + artifact_path, + symbolic_dimensions=self._symbolic_dimensions, + ) with _translate_native_errors("load"): model = runtime.load_model(str(artifact_path)) @@ -969,6 +1017,65 @@ def compile(self) -> None: """Build the Runtime pipeline if it has not already been built.""" self._ensure_built() + def compile_to_file(self, output_path: str | Path, *, validate: bool = True) -> None: + """Persist this CGIR model as device-targeted CG Output IR. + + The output is compiled through the same Runtime compiler path used by + session-backed ONNX-to-CGC execution. Optional validation reloads the + artifact, builds a stage and compares normalized I/O schemas. + """ + if self._backend != "cgc" or not self._is_mlir: + raise ValueError("Persistent Runtime compilation requires a CGIR model.") + destination = Path(output_path) + wr = import_runtime() + runtime = wr.Runtime() + resolved_target = self._resolve_target(runtime, wr) + source_model = compiled_model = builder = stage = pipeline = source_schema = None + try: + with _translate_native_errors("load"): + source_model = runtime.load_model(str(self._model_path)) + source_schema = source_model.schema() + _compile_runtime_model( + source_model, + resolved_target, + destination, + compiled_form=wr.CompiledModelForm.DEVICE_TARGETED, + ) + if not validate: + return + with _translate_native_errors("load"): + compiled_model = runtime.load_model(str(destination)) + with _translate_native_errors("build", device_class=resolved_target.device_class): + builder = runtime.create_pipeline_builder() + stage = builder.add_model_stage( + compiled_model, resolved_target.execution_target + ) + pipeline = builder.build() + compiled_schema = _stage_schema(wr, stage) + try: + if _normalized_schema_descriptors( + source_schema + ) != _normalized_schema_descriptors(compiled_schema): + raise click.ClickException( + "CG Output IR I/O schema differs from input CGIR." + ) + finally: + compiled_schema.close() + finally: + for resource in ( + pipeline, + stage, + builder, + compiled_model, + source_schema, + source_model, + resolved_target.execution_target, + resolved_target.adapter, + runtime, + ): + if resource is not None: + resource.close() + # -- perf --------------------------------------------------------------- @contextmanager def perf(self, warmup: int = 0, monitor: Any | None = None) -> Iterator[PerfContext]: diff --git a/tests/e2e/test_compile_cgc_ir_e2e.py b/tests/e2e/test_compile_cgc_ir_e2e.py new file mode 100644 index 000000000..cf53c1687 --- /dev/null +++ b/tests/e2e/test_compile_cgc_ir_e2e.py @@ -0,0 +1,132 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Native Runtime compilation of generated CGC IR and optional caller input.""" + +import hashlib +import json +import os +import subprocess +import sys +from pathlib import Path + +import onnx +import pytest + +from winml.modelkit.export.cgc import CGCExporter, CGCOptions +from winml.modelkit.sysinfo.dxcore_adapters import enumerate_compute_adapters + + +pytestmark = pytest.mark.e2e +_INPUT_ENCODINGS = ["text", "bytecode"] +if os.environ.get("WINML_TEST_CGC_INPUT_IR"): + _INPUT_ENCODINGS.append("caller") + + +@pytest.fixture +def gpu_luid(): + if sys.platform != "win32": + pytest.skip("Native Windows Runtime GPU compilation requires Windows.") + pytest.importorskip("windowsml.runtime", reason="Preview Runtime API required.") + adapters = [item for item in enumerate_compute_adapters() if item.device_type == "GPU"] + if not adapters: + pytest.skip("No DXCore GPU is available for native compilation.") + return adapters[0].luid + + +@pytest.mark.parametrize("encoding", _INPUT_ENCODINGS) +def test_real_cli_input_ir_to_persistent_output_ir(tmp_path, gpu_luid, encoding): + if encoding == "caller": + source = Path(os.environ["WINML_TEST_CGC_INPUT_IR"]) + assert source.is_file() + else: + source_onnx = tmp_path / "source.onnx" + model = onnx.helper.make_model( + onnx.helper.make_graph( + [onnx.helper.make_node("Relu", ["x"], ["y"])], + "runtime_compile_test", + [onnx.helper.make_tensor_value_info("x", onnx.TensorProto.FLOAT, [2, 4])], + [onnx.helper.make_tensor_value_info("y", onnx.TensorProto.FLOAT, [2, 4])], + ), + opset_imports=[onnx.helper.make_opsetid("", 18)], + ir_version=8, + ) + onnx.save(model, source_onnx) + source = tmp_path / "input.mlir" + CGCExporter(CGCOptions()).export_onnx(source_onnx, source) + if encoding == "bytecode": + import windowsml + from _foundry.compiler import Compiler, SerializationFormat + + with ( + Compiler( + library_path=Path(windowsml.__file__).parent / "lib" / "FoundryToolbox.dll" + ) as compiler, + compiler.create_module_from_mlir(source.read_bytes()) as module, + ): + source.write_bytes(module.serialize(SerializationFormat.BYTECODE)) + original_hash = hashlib.sha256(source.read_bytes()).hexdigest() + output_dir = tmp_path / "compiled" + output = output_dir / "output.mlir" + argv = [ + sys.executable, + "-m", + "winml.modelkit.cli", + "compile", + "-m", + str(source), + "--device", + "gpu", + "--device-luid", + gpu_luid, + "-o", + str(output), + ] + result = subprocess.run( # noqa: S603 -- fixed CLI executable and argv, no shell + argv, + capture_output=True, + text=True, + encoding="utf-8", + timeout=240, + check=False, + ) + assert result.returncode == 0, result.stdout + result.stderr + assert "bytecode" in result.stdout and "schema passed" in result.stdout + assert output.read_bytes().startswith(b"ML\xefR") + receipt = json.loads((output_dir / "output_compile.json").read_text()) + assert receipt["input"]["sha256"] == original_hash + assert receipt["gpu"]["luid"] == gpu_luid + assert receipt["validation"]["status"] == "passed" + assert receipt["compiled_form"] == "DEVICE_TARGETED" + for item in receipt["artifacts"]: + path = output_dir / item["filename"] + assert path.is_file() + assert hashlib.sha256(path.read_bytes()).hexdigest() == item["sha256"] + assert hashlib.sha256(source.read_bytes()).hexdigest() == original_hash + retry = subprocess.run( # noqa: S603 -- same fixed argv for overwrite guard verification + argv, + capture_output=True, + text=True, + encoding="utf-8", + timeout=60, + check=False, + ) + assert retry.returncode != 0 and "--overwrite" in retry.stdout + retry.stderr + if encoding == "text": + invalid = tmp_path / "not-input-ir.onnx" + invalid.write_bytes(source_onnx.read_bytes()) + wrong_args = list(argv) + wrong_args[wrong_args.index("-m", 3) + 1] = str(invalid) + wrong_args[wrong_args.index("-o") + 1] = str(output_dir / "invalid-output.mlir") + wrong = subprocess.run( # noqa: S603 -- fixed CLI argv, no shell + wrong_args, + capture_output=True, + text=True, + encoding="utf-8", + timeout=60, + check=False, + ) + assert wrong.returncode != 0 + assert "Cannot parse CGC Input IR" in wrong.stdout + wrong.stderr + assert not (output_dir / "invalid-output.mlir").exists() diff --git a/tests/unit/commands/test_compile_cgc_ir.py b/tests/unit/commands/test_compile_cgc_ir.py new file mode 100644 index 000000000..0c66b5162 --- /dev/null +++ b/tests/unit/commands/test_compile_cgc_ir.py @@ -0,0 +1,173 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Explicit Input IR / Output IR CLI routing, without native dependencies.""" + +from unittest.mock import Mock + +import pytest +from click.testing import CliRunner + +from winml.modelkit.commands.compile import compile + + +@pytest.fixture +def ir_cli(tmp_path, monkeypatch): + source = tmp_path / "input.mlir" + source.write_bytes(b"input") + operation = Mock(return_value=tmp_path / "output.mlir") + monkeypatch.setattr("winml.modelkit.compiler.compile_cgc_ir", operation) + resolver = Mock(side_effect=AssertionError("IR compile must not resolve an ORT EP")) + monkeypatch.setattr("winml.modelkit.commands.compile.resolve_device", resolver) + args = [ + "-m", + str(source), + ] + return source, operation, args + + +def test_explicit_ir_compilation_routes_without_ep_resolution(ir_cli): + source, operation, args = ir_cli + destination = source.parent / "compiled.mlir" + result = CliRunner().invoke( + compile, + [ + *args, + "--compiler", + "cgc", + "--device", + "gpu", + "--device-luid", + "0x00000000_0x00000001", + "-o", + str(destination), + "--overwrite", + ], + ) + assert result.exit_code == 0, result.output + operation.assert_called_once_with( + source, + destination, + device_luid="0x00000000_0x00000001", + validate=True, + overwrite=True, + ) + assert "bytecode" in result.output and "no inference" in result.output + + +def test_ir_defaults_and_skipped_validation(ir_cli): + source, operation, args = ir_cli + result = CliRunner().invoke(compile, [*args, "--no-validate"]) + assert result.exit_code == 0, result.output + operation.assert_called_once_with( + source, + source.with_name("input_output.mlir"), + device_luid=None, + validate=False, + overwrite=False, + ) + assert "skipped" in result.output + + +def test_ir_output_file_precedes_output_directory(ir_cli): + source, operation, args = ir_cli + destination = source.parent / "selected.mlir" + result = CliRunner().invoke( + compile, + [*args, "-o", str(destination), "--output-dir", str(source.parent / "unused")], + ) + assert result.exit_code == 0, result.output + assert operation.call_args.args[1] == destination + + +def test_explicit_cgc_compiler_routes_nonstandard_suffix(ir_cli): + source, operation, _ = ir_cli + source = source.with_suffix(".cgir") + source.write_bytes(b"input") + result = CliRunner().invoke( + compile, + ["-m", str(source), "--compiler", "cgc"], + ) + assert result.exit_code == 0, result.output + operation.assert_called_once() + + +@pytest.mark.parametrize( + "extra", + [ + ["--ep", "winmlcg"], + ["--ep-options", "device_id=0"], + ["--embed"], + ["--no-embed"], + ["--list"], + ["--device", "cpu"], + ["--device", "npu"], + ["--device", "auto"], + ], +) +def test_ir_rejects_unrelated_options(ir_cli, extra): + _, operation, args = ir_cli + result = CliRunner().invoke(compile, [*args, *extra]) + assert result.exit_code == 2, result.output + operation.assert_not_called() + + +@pytest.mark.parametrize("option", ["--config", "--qnn-sdk-root"]) +def test_ir_rejects_config_and_sdk(ir_cli, option): + source, operation, args = ir_cli + path = source.parent / "config.json" + path.write_text("{}") + result = CliRunner().invoke(compile, [*args, option, str(path)]) + assert result.exit_code == 2, result.output + assert "does not support" in result.output + operation.assert_not_called() + + +def test_ir_rejects_multiple_models(ir_cli): + source, operation, args = ir_cli + result = CliRunner().invoke(compile, [*args, "-m", str(source), "--compiler", "cgc"]) + assert result.exit_code == 2 + assert "exactly one" in result.output + operation.assert_not_called() + + +def test_ir_requires_model(ir_cli): + _, operation, _ = ir_cli + result = CliRunner().invoke(compile, ["--compiler", "cgc"]) + assert result.exit_code == 2 + assert "exactly one" in result.output + operation.assert_not_called() + + +def test_epcontext_rejects_ir_only_device_luid(ir_cli): + source, operation, _ = ir_cli + source = source.with_suffix(".onnx") + source.write_bytes(b"input") + result = CliRunner().invoke( + compile, + ["-m", str(source), "--device-luid", "0x00000000_0x00000001"], + ) + assert result.exit_code == 2 + assert "only supported for CG Output IR" in result.output + operation.assert_not_called() + + +def test_ir_filesystem_error_is_actionable(ir_cli): + _, operation, args = ir_cli + operation.side_effect = OSError("Disk write failed") + result = CliRunner().invoke(compile, args) + assert result.exit_code == 1 + assert "CGC compilation failed: Disk write failed" in result.output + + +def test_ir_flags_and_encoding_visible_in_help(): + result = CliRunner().invoke(compile, ["--help"]) + assert result.exit_code == 0 + for value in ( + "--compiler", + "cgc", + ".mlir", + "bytecode", + ): + assert value in result.output diff --git a/tests/unit/compiler/test_runtime_compile.py b/tests/unit/compiler/test_runtime_compile.py new file mode 100644 index 000000000..21833098d --- /dev/null +++ b/tests/unit/compiler/test_runtime_compile.py @@ -0,0 +1,243 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Persistence and publication tests for Runtime-backed CGIR compilation.""" + +import json +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import Mock + +import click +import pytest + +from winml.modelkit.compiler import compile_cgc_ir +from winml.modelkit.compiler.runtime import _publish, _select_gpu + + +@pytest.fixture +def native(tmp_path, monkeypatch): + source = tmp_path / "input.mlir" + source.write_bytes(b"Input IR fixture") + output = tmp_path / "output.mlir" + source.with_name("input_metadata.json").write_text( + json.dumps({"inputs": [{"name": "x", "index": 0}], "outputs": [{"name": "y", "index": 0}]}), + ) + compile_action = Mock() + + def write_artifact(path: Path) -> None: + path.write_bytes(b"ML\xefRcompiled") + path.with_suffix(".data").write_bytes(b"resources") + + compile_action.side_effect = write_artifact + sessions = [] + + class FakeSession: + def __init__(self, model_path, *, backend, device_luid): + self.model_path = Path(model_path) + self.backend = backend + self.device_luid = device_luid + self.validate = None + sessions.append(self) + + def compile_to_file(self, path, *, validate): + self.validate = validate + compile_action(Path(path)) + + monkeypatch.setattr("winml.modelkit.compiler.runtime.WinMLRuntimeSession", FakeSession) + gpu = SimpleNamespace(device_type="GPU", name="Test GPU", luid="0x00000000_0x00000001") + monkeypatch.setattr("winml.modelkit.compiler.runtime.enumerate_compute_adapters", lambda: [gpu]) + return SimpleNamespace( + source=source, + output=output, + compile_action=compile_action, + sessions=sessions, + ) + + +def test_compile_routes_through_runtime_session_and_publishes_bundle(native): + result = compile_cgc_ir(native.source, native.output) + + assert result == native.output + assert native.output.read_bytes() == b"ML\xefRcompiled" + assert native.output.with_suffix(".data").read_bytes() == b"resources" + metadata = json.loads(native.output.with_name("output_metadata.json").read_text()) + assert metadata["format"] == "cg-output-ir" + receipt = json.loads(native.output.with_name("output_compile.json").read_text()) + assert receipt["validation"]["status"] == "passed" + assert receipt["output_encoding"] == "mlir-bytecode" + assert {entry["filename"] for entry in receipt["artifacts"]} == { + "output.mlir", + "output.data", + "output_metadata.json", + } + session = native.sessions[0] + assert session.model_path == native.source + assert session.backend == "cgc" + assert session.device_luid == 1 + assert session.validate is True + + +def test_no_validate_is_forwarded_to_runtime_session(native): + compile_cgc_ir(native.source, native.output, validate=False) + assert native.sessions[0].validate is False + receipt = json.loads(native.output.with_name("output_compile.json").read_text()) + assert receipt["validation"]["status"] == "skipped" + + +def test_overwrite_requires_opt_in_and_preserves_source(native): + native.output.write_bytes(b"existing") + with pytest.raises(click.ClickException, match="already exists"): + compile_cgc_ir(native.source, native.output) + native.compile_action.assert_not_called() + + compile_cgc_ir(native.source, native.output, overwrite=True) + assert native.output.read_bytes() == b"ML\xefRcompiled" + assert native.source.read_bytes() == b"Input IR fixture" + + +def test_overwrite_removes_stale_output_metadata(native): + native.source.with_name("input_metadata.json").unlink() + native.output.write_bytes(b"old model") + stale = native.output.with_name("output_metadata.json") + stale.write_text('{"inputs": [{"name": "old", "index": 0}]}') + + compile_cgc_ir(native.source, native.output, overwrite=True) + + assert native.output.read_bytes() == b"ML\xefRcompiled" + assert not stale.exists() + + +def test_companion_collision_preserves_prior_output(native): + sidecar = native.output.with_suffix(".data") + sidecar.write_bytes(b"prior sidecar") + with pytest.raises(click.ClickException, match="already exists"): + compile_cgc_ir(native.source, native.output) + assert sidecar.read_bytes() == b"prior sidecar" + assert not native.output.exists() + + +def test_compile_failure_does_not_replace_existing_output(native): + native.output.write_bytes(b"old") + native.compile_action.side_effect = RuntimeError("compiler error") + with pytest.raises(RuntimeError, match="compiler error"): + compile_cgc_ir(native.source, native.output, overwrite=True) + assert native.output.read_bytes() == b"old" + + +@pytest.mark.parametrize("failure", ["empty", "text"]) +def test_rejects_invalid_runtime_output(native, failure): + if failure == "empty": + native.compile_action.side_effect = lambda path: None + else: + native.compile_action.side_effect = lambda path: path.write_text("unexpected") + with pytest.raises(click.ClickException): + compile_cgc_ir(native.source, native.output) + assert not native.output.exists() + + +def test_changed_input_is_rejected_before_publication(native): + def mutate(path): + path.write_bytes(b"ML\xefRcompiled") + native.source.write_bytes(b"changed") + + native.compile_action.side_effect = mutate + with pytest.raises(click.ClickException, match="changed during compilation"): + compile_cgc_ir(native.source, native.output) + assert not native.output.exists() + + +@pytest.mark.parametrize("kind", ["invalid-json", "list"]) +def test_invalid_metadata_is_explicit(native, kind): + native.source.with_name("input_metadata.json").write_text( + "{" if kind == "invalid-json" else "[]" + ) + with pytest.raises(click.ClickException, match="metadata"): + compile_cgc_ir(native.source, native.output) + + +def test_source_metadata_change_during_compile_is_rejected(native): + def mutate(path): + path.write_bytes(b"ML\xefRcompiled") + native.source.with_name("input_metadata.json").write_text("{}") + + native.compile_action.side_effect = mutate + with pytest.raises(click.ClickException, match="metadata changed"): + compile_cgc_ir(native.source, native.output) + assert not native.output.exists() + + +def test_input_and_output_metadata_collision_preserves_source(native): + destination = native.source.with_suffix(".mlirbc") + with pytest.raises(click.ClickException, match="source artifact"): + compile_cgc_ir(native.source, destination, overwrite=True) + assert not destination.exists() + + +def test_nonexistent_input_and_wrong_output_suffix_are_explicit(native): + with pytest.raises(click.ClickException, match="file not found"): + compile_cgc_ir(native.source.with_name("missing.mlir"), native.output) + with pytest.raises(click.UsageError, match=r"\.mlir or \.mlirbc"): + compile_cgc_ir(native.source, native.output.with_suffix(".onnx")) + native.compile_action.assert_not_called() + + +def test_in_place_compile_is_rejected_even_with_overwrite(native): + with pytest.raises(click.ClickException, match="different files"): + compile_cgc_ir(native.source, native.source, overwrite=True) + native.compile_action.assert_not_called() + + +def test_cgc_compile_without_io_metadata(native): + native.source.with_name("input_metadata.json").unlink() + compile_cgc_ir(native.source, native.output) + assert native.output.is_file() + assert not native.output.with_name("output_metadata.json").exists() + + +@pytest.mark.parametrize("count", [0, 2]) +def test_gpu_selection_rejects_missing_or_ambiguous_adapters(monkeypatch, count): + adapters = [ + SimpleNamespace(device_type="GPU", name=f"GPU {i}", luid=f"0x00000000_0x0000000{i}") + for i in range(count) + ] + monkeypatch.setattr( + "winml.modelkit.compiler.runtime.enumerate_compute_adapters", lambda: adapters + ) + with pytest.raises(click.ClickException, match="unambiguous"): + _select_gpu(None) + if count == 2: + assert _select_gpu("0X00000000_0X00000001") == adapters[1] + + +def test_atomic_publication_rolls_back_new_and_invalidated_files(tmp_path, monkeypatch): + staging = tmp_path / "stage" + staging.mkdir() + (staging / "first").write_bytes(b"new first") + (staging / "second").write_bytes(b"new second") + (tmp_path / "first").write_bytes(b"old first") + (tmp_path / "second").write_bytes(b"old second") + stale = tmp_path / "stale" + stale.write_bytes(b"old metadata") + replace = Path.replace + + def fail_one_publish(self, target): + if self == staging / "second": + raise OSError("publish failure") + return replace(self, target) + + monkeypatch.setattr(Path, "replace", fail_one_publish) + with pytest.raises(OSError, match="publish failure"): + _publish(staging, tmp_path / "first", True, (), invalidated=(stale,)) + assert (tmp_path / "first").read_bytes() == b"old first" + assert (tmp_path / "second").read_bytes() == b"old second" + assert stale.read_bytes() == b"old metadata" + + +def test_nested_runtime_layout_is_not_silently_discarded(tmp_path): + staging = tmp_path / "stage" + (staging / "resources").mkdir(parents=True) + (staging / "compiled.mlir").write_bytes(b"ML\xefRcompiled") + with pytest.raises(click.ClickException, match="nested Runtime"): + _publish(staging, tmp_path / "compiled.mlir", False, ()) diff --git a/tests/unit/export/cgc/test_exporter.py b/tests/unit/export/cgc/test_exporter.py index deffbf968..32db74ae3 100644 --- a/tests/unit/export/cgc/test_exporter.py +++ b/tests/unit/export/cgc/test_exporter.py @@ -3,6 +3,7 @@ # Licensed under the MIT License. # -------------------------------------------------------------------------- +import json import sys from enum import Enum from pathlib import Path @@ -204,7 +205,7 @@ def test_external_mlir_overwrite_replaces_sidecar(tmp_path: Path) -> None: assert output.read_text(encoding="utf-8") == "module { cgc.test }" assert sidecar.read_bytes() == b"new weights" - assert metadata.is_file() + assert json.loads(metadata.read_text(encoding="utf-8"))["format"] == "cgc-input-ir" def test_external_mlir_failure_preserves_existing_bundle(tmp_path: Path) -> None: diff --git a/tests/unit/session/test_runtime_session.py b/tests/unit/session/test_runtime_session.py index f5ad8b673..02702a4d7 100644 --- a/tests/unit/session/test_runtime_session.py +++ b/tests/unit/session/test_runtime_session.py @@ -9,8 +9,9 @@ import json import threading import time +from pathlib import Path from types import SimpleNamespace -from typing import TYPE_CHECKING, Any +from typing import Any from unittest.mock import Mock import click @@ -24,14 +25,11 @@ _numpy_dtype_for, _ResolvedRuntimeTarget, _shape_with_dynamic_dims, + _stage_schema, _symbolic_dimensions_for_inputs, ) -if TYPE_CHECKING: - from pathlib import Path - - class _NotSupportedError(Exception): pass @@ -261,6 +259,70 @@ def test_schema_helpers() -> None: None, 4, ] + + +def test_stage_schema_falls_back_to_projection_interface() -> None: + interface = Mock() + wrapped = Mock() + stage = SimpleNamespace(interface=Mock()) + wr = SimpleNamespace( + _b=SimpleNamespace(IWinMLStageSchema=interface), + ModelSchema=Mock(return_value=wrapped), + ) + + assert _stage_schema(wr, stage) is wrapped + stage.interface.QueryInterface.assert_called_once_with(interface) + wr.ModelSchema.assert_called_once_with( + stage.interface.QueryInterface.return_value, stage + ) + + +def test_persistent_compile_normalizes_dynamic_schema_dimensions(tmp_path, monkeypatch) -> None: + source_schema = Mock() + source_schema.inputs.return_value = [("float32", (0, 4))] + source_schema.outputs.return_value = [("float32", (0, 4))] + compiled_schema = Mock() + compiled_schema.inputs.return_value = [("float32", ((1 << 64) - 1, 4))] + compiled_schema.outputs.return_value = [("float32", ((1 << 64) - 1, 4))] + source_model = Mock(schema=Mock(return_value=source_schema)) + compiled_model = Mock() + options = SimpleNamespace(form=None, close=Mock()) + compiler = Mock( + supports_form=Mock(return_value=True), + create_options=Mock(return_value=options), + ) + output = tmp_path / "output.mlir" + compiler.compile_to_file.side_effect = lambda _model, path, **_kwargs: Path( + path + ).write_bytes(b"ML\xefRcompiled") + target = Mock(model_compiler=Mock(return_value=compiler)) + resolved = _ResolvedRuntimeTarget(target, "gpu") + stage = Mock(schema=Mock(return_value=compiled_schema)) + builder = Mock() + builder.add_model_stage.return_value = stage + runtime = Mock() + runtime.load_model.side_effect = [source_model, compiled_model] + runtime.create_pipeline_builder.return_value = builder + wr = SimpleNamespace( + Runtime=Mock(return_value=runtime), + CompiledModelForm=SimpleNamespace(DEVICE_TARGETED="device-targeted"), + ) + monkeypatch.setattr( + "winml.modelkit.session.runtime_session.import_runtime", lambda: wr + ) + session = WinMLRuntimeSession( + tmp_path / "input.mlir", backend="cgc", device_luid=1 + ) + monkeypatch.setattr(session, "_resolve_target", lambda _runtime, _wr: resolved) + + session.compile_to_file(output) + + assert output.is_file() + assert options.form == "device-targeted" + compiler.compile_to_file.assert_called_once_with( + source_model, str(output), options=options + ) + compiled_schema.close.assert_called_once_with() with pytest.raises(click.ClickException): _numpy_dtype_for(SimpleNamespace(name="BFLOAT16")) @@ -523,8 +585,8 @@ def test_prepare_inputs_preserves_shape( session.close() -def test_mlir_session_requires_ep_device() -> None: - with pytest.raises(ValueError, match="ep_device is required"): +def test_mlir_session_requires_target_device() -> None: + with pytest.raises(ValueError, match="ep_device or device_luid is required"): WinMLRuntimeSession("model.mlir", backend="cgc")