Add EP and device targeting to optimize

src/winml/modelkit/commands/optimize.py

@@ -229,7 +229,26 @@ def _render_check_optim(console: Console, findings: list[Any], verbose: bool) ->
     )
 
 
-def _run_check_optim(model: Path, all_caps: dict[str, Any], verbose: bool) -> None:
+def _resolve_optimization_target(ep: str | None, device: str | None) -> tuple[Any, Any]:
+    """Resolve an explicit optimization EP/device request."""
+    from ..session import (
+        EPDeviceTarget,
+        WinMLEPRegistry,
+        resolve_device,
+    )
+
+    target = resolve_device(EPDeviceTarget(ep=ep or "auto", device=(device or "auto").lower()))
+    return target, WinMLEPRegistry.instance().auto_device(target)
+
+
+def _run_check_optim(
+    model: Path,
+    all_caps: dict[str, Any],
+    verbose: bool,
+    *,
+    target: Any | None = None,
+    ep_device: Any | None = None,
+) -> None:
     """Probe which optimizations apply to the model and print a report.
 
     No output file is written. Every boolean capability that is off by default
@@ -239,6 +258,8 @@ def _run_check_optim(model: Path, all_caps: dict[str, Any], verbose: bool) -> No
         model: Path to the input ONNX model.
         all_caps: The full capability registry.
         verbose: Whether to show every affected node/constant.
+        target: Resolved EP/device target, if explicitly requested.
+        ep_device: Resolved runtime EP device used by each optimization probe.
     """
     from ..optim import BoolCapability, analyze_model
 
@@ -247,6 +268,8 @@ def _run_check_optim(model: Path, all_caps: dict[str, Any], verbose: bool) -> No
     )
 
     console.print(f"[bold blue]Input:[/bold blue] {model}")
+    if target is not None:
+        console.print(f"[bold blue]Target:[/bold blue] {target.ep} on {target.device.upper()}")
     console.print(
         "[dim]--check-optim — analyzing applicable optimizations (no output written).[/dim]"
     )
@@ -268,6 +291,7 @@ def _run_check_optim(model: Path, all_caps: dict[str, Any], verbose: bool) -> No
         findings = analyze_model(
             onnx_model,
             all_caps,
+            ep_device=ep_device,
             on_probe_start=lambda name: progress.update(task_id, description=name),
             on_probe_complete=lambda _: progress.advance(task_id),
         )
@@ -303,6 +327,19 @@ def _run_check_optim(model: Path, all_caps: dict[str, Any], verbose: bool) -> No
 )
 @cli_utils.output_option("Output path (default: {input}_opt.onnx)")
 @cli_utils.overwrite_option()
+@cli_utils.ep_option(
+    required=False,
+    default=None,
+    include_auto=True,
+    include_cuda=True,
+    optional_message="If omitted with --device, selects a compatible EP automatically.",
+)
+@cli_utils.device_option(
+    required=False,
+    default=None,
+    include_auto=True,
+    optional_message="If both --ep and --device are omitted, preserves CPU optimization.",
+)
 @click.option(
     "--config",
     "-c",
@@ -322,6 +359,8 @@ def optimize(
     model: Path | None,
     output: Path | None,
     overwrite: bool,
+    ep: str | None,
+    device: str | None,
     config: Path | None,
     verbose: int,
     quiet: bool,
@@ -476,9 +515,37 @@ def optimize(
     verbose, quiet = cli_utils.resolve_verbosity(ctx, verbose, quiet)
     configure_logging(verbosity=verbose, quiet=quiet)
 
+    target = None
+    ep_device = None
+    if ep is not None or device is not None:
+        from ..session import (
+            DeviceNotFound,
+            UnknownListingPick,
+            WinMLEPNotDiscovered,
+            WinMLEPRegistrationFailed,
+        )
+
+        try:
+            target, ep_device = _resolve_optimization_target(ep, device)
+        except (
+            DeviceNotFound,
+            RuntimeError,
+            UnknownListingPick,
+            ValueError,
+            WinMLEPNotDiscovered,
+            WinMLEPRegistrationFailed,
+        ) as e:
+            raise click.UsageError(f"Could not resolve optimization target: {e}") from e
+
     # Handle --check-optim: report which optimizations apply, write nothing.
     if check_optim:
-        _run_check_optim(model, all_caps, bool(verbose))
+        _run_check_optim(
+            model,
+            all_caps,
+            bool(verbose),
+            target=target,
+            ep_device=ep_device,
+        )
         return
 
     # Import optimizer
@@ -542,6 +609,14 @@ def optimize(
     if verbose:
         optimizer_kwargs["verbose"] = True
 
+    if ep_device is not None:
+        if target is None:
+            raise RuntimeError("Resolved optimization EP device has no target metadata.")
+        optimizer_kwargs["ep_device"] = ep_device
+        console.print(
+            f"[bold blue]Target:[/bold blue] {target.ep} on {target.device.upper()}"
+        )
+
     try:
         console.print("\n[bold]Loading model...[/bold]")
         onnx_model = load_onnx(model)
@@ -556,10 +631,22 @@ def optimize(
 
         # Report results
         optimized_nodes = len(optimized_model.graph.node)
-        reduction = (1 - optimized_nodes / original_nodes) * 100 if original_nodes else 0
+        node_change = (
+            abs(optimized_nodes - original_nodes) / original_nodes * 100 if original_nodes else 0
+        )
 
         console.print(f"\n[bold green]Success![/bold green] Model optimized: {output}")
-        node_info = f"Nodes: {original_nodes} -> {optimized_nodes} ({reduction:.1f}% reduction)"
+        if optimized_nodes < original_nodes:
+            change_label = "reduction"
+        elif optimized_nodes > original_nodes:
+            change_label = "increase"
+        else:
+            change_label = None
+        node_info = f"Nodes: {original_nodes} -> {optimized_nodes}"
+        if change_label is None:
+            node_info += " (no change)"
+        else:
+            node_info += f" ({node_change:.1f}% {change_label})"
         console.print(f"[dim]{node_info}[/dim]")
 
     except Exception as e:

src/winml/modelkit/optim/analysis.py

@@ -283,8 +283,10 @@ def _run_pipe(pipe: Any, model: ModelProto, config: Any) -> ModelProto:
 def _iter_findings(
     model: ModelProto,
     capabilities: dict[str, CapabilityDef],
+    *,
     on_probe_start: Callable[[str], None] | None = None,
     on_probe_complete: Callable[[str], None] | None = None,
+    **optimizer_kwargs: Any,
 ) -> Iterator[tuple[CapabilityFinding, ModelProto]]:
     """Yield ``(finding, produced_model)`` for every applicable optimization.
 
@@ -319,6 +321,7 @@ def complete_probe(cap_name: str) -> None:
 
     # Baseline kwargs = every capability at its default value.
     default_kwargs = {cap.python_name: cap.default for cap in capabilities.values()}
+    default_kwargs.update(optimizer_kwargs)
     kebab_defaults = {name: cap.default for name, cap in capabilities.items()}
 
     # Mandatory pre-stage — mirrors Optimizer.optimize(). The clone is required:
@@ -358,6 +361,22 @@ def complete_probe(cap_name: str) -> None:
             if on_probe_start is not None:
                 on_probe_start(cap_name)
             try:
+                ep_device = optimizer_kwargs.get("ep_device")
+                if ep_device is not None and cap.ep_constraint is not None:
+                    from ..utils.constants import normalize_ep_name
+
+                    target_ep = normalize_ep_name(ep_device.device.ep_name)
+                    if not any(
+                        normalize_ep_name(name) == target_ep for name in cap.ep_constraint
+                    ):
+                        logger.debug(
+                            "Skipping capability '%s': target EP %s is not in %s",
+                            cap.name,
+                            target_ep,
+                            cap.ep_constraint,
+                        )
+                        continue
+
                 # Enable only this capability (plus its dependencies) on top of
                 # the all-defaults configuration.
                 kebab = dict(kebab_defaults)
@@ -368,6 +387,7 @@ def complete_probe(cap_name: str) -> None:
                     for name, value in kebab.items()
                     if name in capabilities
                 }
+                probe_kwargs.update(optimizer_kwargs)
 
                 probe_config = pipe.build_config(**probe_kwargs)
                 should_process = getattr(pipe, "should_process", None)
@@ -428,6 +448,7 @@ def analyze_model(
     *,
     on_probe_start: Callable[[str], None] | None = None,
     on_probe_complete: Callable[[str], None] | None = None,
+    **optimizer_kwargs: Any,
 ) -> list[CapabilityFinding]:
     """Probe every applicable optimization capability against ``model``.
 
@@ -440,6 +461,8 @@ def analyze_model(
         model: The input ONNX model (never modified).
         capabilities: The full capability registry (kebab-case keyed), e.g.
             from ``optim.pipes.get_all_capabilities()``.
+        **optimizer_kwargs: Pipeline context forwarded to every pipe configuration,
+            such as a resolved ``ep_device``.
         on_probe_start: Optional callback invoked with the capability name
             before each probe begins.
         on_probe_complete: Optional callback invoked with the capability name
@@ -456,13 +479,15 @@ def analyze_model(
             capabilities,
             on_probe_start=on_probe_start,
             on_probe_complete=on_probe_complete,
+            **optimizer_kwargs,
         )
     ]
 
 
 def iter_optimization_outputs(
     model: ModelProto,
     capabilities: dict[str, CapabilityDef],
+    **optimizer_kwargs: Any,
 ) -> Iterator[tuple[CapabilityFinding, ModelProto]]:
     """Yield each applicable optimization together with the model it produces.
 
@@ -476,10 +501,11 @@ def iter_optimization_outputs(
     Args:
         model: The input ONNX model (never modified).
         capabilities: The full capability registry (kebab-case keyed).
+        **optimizer_kwargs: Pipeline context forwarded to every pipe configuration.
 
     Yields:
         ``(finding, produced_model)`` pairs in pipeline order, one per applicable
         optimization. The pairs are produced lazily; materialize the iterator if
         the produced models must outlive iteration.
     """
-    yield from _iter_findings(model, capabilities)
+    yield from _iter_findings(model, capabilities, **optimizer_kwargs)

src/winml/modelkit/optim/pipes/graph.py

@@ -26,6 +26,8 @@
 if TYPE_CHECKING:
     import onnx
 
+    from ...session import WinMLEPDevice
+
 # Import all capability modules to build capabilities dict
 from ..capabilities import (
     activation,
@@ -146,19 +148,23 @@ def __init__(
         self,
         enabled: list[str] | None = None,
         verbose: bool = False,
+        ep_device: WinMLEPDevice | None = None,
     ) -> None:
         """Initialize with all advanced optimizers disabled, enable specified ones.
 
         Args:
             enabled: List of python_names to enable (e.g., ["gelu_fusion"])
             verbose: Enable verbose logging
+            ep_device: Resolved EP/device target for provider-aware optimization.
+                If omitted, optimization uses CPUExecutionProvider.
 
         Note:
             Only default=False capabilities are managed here. Basic ORT
             optimizations (ConstantFolding, etc.) run automatically at Level 2.
         """
         self.optimization_level = 2  # Finalized at Level 2
         self.verbose = verbose
+        self.ep_device = ep_device
 
         # Special flag for GeluApproximation (requires separate session config)
         # ORT docs: "GeluApproximation has side effects which may change results.
@@ -356,7 +362,11 @@ def build_config(cls, **kwargs: Any) -> ORTGraphPipeConfig:
                 and cap.python_name not in explicitly_disabled
             ]
 
-        config = ORTGraphPipeConfig(enabled=enabled, verbose=verbose)
+        config = ORTGraphPipeConfig(
+            enabled=enabled,
+            verbose=verbose,
+            ep_device=kwargs.get("ep_device"),
+        )
 
         # Explicitly disable capabilities that user set to False
         # This handles default=True caps like constant_folding
@@ -456,7 +466,12 @@ def _log_process_verbose(
             "  graph_optimization_level: %d (ORT_ENABLE_EXTENDED)", config.optimization_level
         )
         logger.debug("  optimized_model_filepath: %s", output_file)
-        logger.debug("  providers: ['CPUExecutionProvider']")
+        provider = (
+            config.ep_device.device.ep_name
+            if config.ep_device is not None
+            else "CPUExecutionProvider"
+        )
+        logger.debug("  provider: %s", provider)
 
         # Session config entries
         logger.debug("[Session Config Entries]")
@@ -562,15 +577,38 @@ def process(self, model: onnx.ModelProto, config: ORTGraphPipeConfig) -> onnx.Mo
                     "1",
                 )
 
+            if config.ep_device is not None:
+                from ...session import lookup_device_spec
+
+                spec = lookup_device_spec(
+                    config.ep_device.device.ep_name,
+                    config.ep_device.device.device_type.lower(),
+                )
+                if spec is None:
+                    logger.debug(
+                        "No device specification found for %s on %s; "
+                        "using empty provider options",
+                        config.ep_device.device.ep_name,
+                        config.ep_device.device.device_type,
+                    )
+                provider_options = dict(spec.default_provider_options) if spec else {}
+                sess_opts.add_provider_for_devices(
+                    [config.ep_device.device.ort_handle],
+                    provider_options,
+                )
+
             # Verbose output for process
             if config.verbose:
                 self._log_process_verbose(config, model, input_file, output_file, disable_list)
 
             # Create session to trigger optimization
             try:
-                _ = ort.InferenceSession(
-                    str(input_file), sess_opts, providers=["CPUExecutionProvider"]
-                )
+                if config.ep_device is None:
+                    session = ort.InferenceSession(
+                        str(input_file), sess_opts, providers=["CPUExecutionProvider"]
+                    )
+                else:
+                    session = ort.InferenceSession(str(input_file), sess_opts)
             except Exception as e:
                 raise OptimizationError(
                     f"ONNX Runtime optimization failed: {e}",
@@ -582,6 +620,22 @@ def process(self, model: onnx.ModelProto, config: ORTGraphPipeConfig) -> onnx.Mo
                     cause=e,
                 ) from e
 
+            if config.ep_device is not None:
+                from ...utils.constants import normalize_ep_name
+
+                expected_provider = normalize_ep_name(config.ep_device.device.ep_name)
+                active_providers = session.get_providers()
+                logger.debug("ORT optimization session providers: %s", active_providers)
+                if not any(
+                    normalize_ep_name(provider) == expected_provider
+                    for provider in active_providers
+                ):
+                    raise OptimizationError(
+                        f"Requested provider {expected_provider} was not activated; "
+                        f"active providers: {active_providers}",
+                        pipe_name=self.name,
+                    )
+
             # Load and return optimized model
             try:
                 return load_onnx(output_file, validate=False)