Add VAE graph optimization surgeries

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/capabilities/surgery.py

@@ -5,7 +5,8 @@
 """Surgery capabilities for precise model modifications.
 
 These capabilities perform targeted graph transformations that are not part of
-ONNX Runtime's standard optimization passes. They run before ORT optimizations.
+ONNX Runtime's standard optimization passes. Each runs at the pipeline stage
+where its required graph evidence is available.
 
 Use cases:
 - Fix quantization issues (extreme values, invalid scales)
@@ -54,3 +55,44 @@
     category=CapabilityCategory.SURGERY,
     default=False,
 )
+
+# Remove operations that are provably redundant around channel-wise L2
+# normalization while preserving the original ReduceL2 and Div kernels.
+SIMPLIFY_L2_NORMALIZATION = BoolCapability(
+    name="simplify-l2-normalization",
+    ort_name=None,
+    description="Remove redundant Clip(min=0) and Expand from ReduceL2-based normalization",
+    category=CapabilityCategory.SURGERY,
+    default=False,
+)
+
+# Replace the exporter-generated integer grid used for 2x nearest-neighbor
+# spatial upsampling with the equivalent standard ONNX Resize operation.
+GATHERND_TO_RESIZE = BoolCapability(
+    name="gathernd-to-resize",
+    ort_name=None,
+    description="Replace exact 2x nearest-neighbor GatherND upsampling grids with Resize",
+    category=CapabilityCategory.SURGERY,
+    default=False,
+)
+
+# Emit the CUDA-supported Microsoft contrib operation for exact SiLU topology.
+SILU_TO_QUICK_GELU = BoolCapability(
+    name="silu-to-quick-gelu",
+    ort_name=None,
+    description="Replace x*Sigmoid(x) with com.microsoft QuickGelu(alpha=1)",
+    category=CapabilityCategory.SURGERY,
+    default=False,
+    ep_constraint=("CUDA",),
+)
+
+# Move a scalar multiplication into the CUDA-supported Microsoft contrib
+# MatMul kernel. This is algebraically equivalent but may change FP rounding.
+SCALED_MATMUL_TO_FUSED_MATMUL = BoolCapability(
+    name="scaled-matmul-to-fused-matmul",
+    ort_name=None,
+    description="Replace MatMul followed by scalar Mul with com.microsoft FusedMatMul",
+    category=CapabilityCategory.SURGERY,
+    default=False,
+    ep_constraint=("CUDA",),
+)

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

@@ -4,9 +4,8 @@
 # --------------------------------------------------------------------------
 """Optimization pipes for ONNX models.
 
-Note: Shape inference (symbolic and standard) are now mandatory stages
-in the Optimizer class, not configurable pipes. Only ORTGraphPipe and
-ORTFusionPipe remain as capability-driven optimization pipes.
+Shape inference is a mandatory Optimizer stage. Capability-driven pipes run
+before, during, and after ORT graph optimization according to their needs.
 """
 
 from typing import Any
@@ -20,20 +19,27 @@
 from .fusion import ORTFusionPipe, ORTFusionPipeConfig
 from .graph import GRAPH_CAPABILITIES, ORTGraphPipe, ORTGraphPipeConfig
 from .rewrite import RewritePipe, RewritePipeConfig
-from .surgery import SURGERY_CAPABILITIES, SurgeryPipe, SurgeryPipeConfig
+from .surgery import (
+    PRE_SURGERY_CAPABILITIES,
+    SURGERY_CAPABILITIES,
+    PreSurgeryPipe,
+    SurgeryPipe,
+    SurgeryPipeConfig,
+)
 
 
 # Optimization pipes to run in sequence
+# - PreSurgeryPipe: Proof-dependent graph rewrites that must inspect the exported graph.
 # - ORTGraphPipe: ORT graph-level optimizations (C++ optimizer), including constant folding.
-#   Runs first so downstream pipes see a constant-folded graph (e.g. Reshape shape inputs
-#   become literal constants, enabling skeleton-based pattern matching).
+#   Runs before downstream pattern pipes so they see a constant-folded graph.
 # - AlgebraicRewritePipe: Exact topology-based algebraic rewrites (after ORT folding).
 # - RewritePipe: Pattern-based subgraph rewriting (runs after ORT constant folding so that
 #   shape constants are visible, but before ORTFusionPipe so normalised patterns are
 #   available for transformer fusions).
 # - ORTFusionPipe: ORT transformer fusions (Python optimizer)
 # - SurgeryPipe: Post-optimization model surgery (runs last to clamp constants after folding)
 PIPES: list[type[BasePipe]] = [
+    PreSurgeryPipe,
     ORTGraphPipe,
     AlgebraicRewritePipe,
     RewritePipe,
@@ -58,6 +64,7 @@ def get_all_capabilities() -> dict[str, Any]:
     "ALGEBRAIC_CAPABILITIES",
     "GRAPH_CAPABILITIES",
     "PIPES",
+    "PRE_SURGERY_CAPABILITIES",
     "SURGERY_CAPABILITIES",
     "AlgebraicRewritePipe",
     "AlgebraicRewritePipeConfig",
@@ -68,6 +75,7 @@ def get_all_capabilities() -> dict[str, Any]:
     "ORTGraphPipeConfig",
     "OptimizationError",
     "PipeConfig",
+    "PreSurgeryPipe",
     "RewritePipe",
     "RewritePipeConfig",
     "SurgeryPipe",