From c04228c25abe8b0113bb8962e70baf850ca1b750 Mon Sep 17 00:00:00 2001 From: hualxie Date: Thu, 30 Jul 2026 16:57:22 +0800 Subject: [PATCH 1/4] Add EP device targeting to optimize --- src/winml/modelkit/commands/optimize.py | 59 +++++++++++++++++++- src/winml/modelkit/optim/pipes/graph.py | 41 ++++++++++++-- tests/unit/commands/test_optimize_cli.py | 64 +++++++++++++++++++++- tests/unit/optim/pipes/test_pipe_config.py | 39 +++++++++++++ 4 files changed, 195 insertions(+), 8 deletions(-) diff --git a/src/winml/modelkit/commands/optimize.py b/src/winml/modelkit/commands/optimize.py index 0b5a7917b..297177200 100644 --- a/src/winml/modelkit/commands/optimize.py +++ b/src/winml/modelkit/commands/optimize.py @@ -290,6 +290,18 @@ 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, + 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", @@ -309,6 +321,8 @@ def optimize( model: Path | None, output: Path | None, overwrite: bool, + ep: str | None, + device: str | None, config: Path | None, verbose: int, quiet: bool, @@ -529,6 +543,36 @@ def optimize( if verbose: optimizer_kwargs["verbose"] = True + if ep is not None or device is not None: + from ..session import ( + DeviceNotFound, + EPDeviceTarget, + UnknownListingPick, + WinMLEPNotDiscovered, + WinMLEPRegistrationFailed, + WinMLEPRegistry, + resolve_device, + ) + + try: + target = resolve_device( + EPDeviceTarget(ep=ep or "auto", device=(device or "auto").lower()) + ) + ep_device = WinMLEPRegistry.instance().auto_device(target) + except ( + DeviceNotFound, + RuntimeError, + UnknownListingPick, + ValueError, + WinMLEPNotDiscovered, + WinMLEPRegistrationFailed, + ) as e: + raise click.UsageError(f"Could not resolve optimization target: {e}") from e + 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) @@ -543,10 +587,21 @@ 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 = "change" + node_info = ( + f"Nodes: {original_nodes} -> {optimized_nodes} " + f"({node_change:.1f}% {change_label})" + ) console.print(f"[dim]{node_info}[/dim]") except Exception as e: diff --git a/src/winml/modelkit/optim/pipes/graph.py b/src/winml/modelkit/optim/pipes/graph.py index 853cdb9ee..6ea5feeae 100644 --- a/src/winml/modelkit/optim/pipes/graph.py +++ b/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,12 +148,15 @@ 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 @@ -159,6 +164,7 @@ def __init__( """ 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,31 @@ 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(), + ) + 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: + _ = ort.InferenceSession( + str(input_file), sess_opts, providers=["CPUExecutionProvider"] + ) + else: + _ = ort.InferenceSession(str(input_file), sess_opts) except Exception as e: raise OptimizationError( f"ONNX Runtime optimization failed: {e}", diff --git a/tests/unit/commands/test_optimize_cli.py b/tests/unit/commands/test_optimize_cli.py index d5b6507aa..58586c1a2 100644 --- a/tests/unit/commands/test_optimize_cli.py +++ b/tests/unit/commands/test_optimize_cli.py @@ -73,7 +73,19 @@ def test_help_exits_cleanly(self, runner: CliRunner) -> None: def test_help_shows_required_flags(self, runner: CliRunner) -> None: result = runner.invoke(optimize, ["--help"]) assert result.exit_code == 0 - for flag in ("--model", "-m", "--output", "-o", "--config", "-c", "--verbose", "-v"): + for flag in ( + "--model", + "-m", + "--output", + "-o", + "--ep", + "--device", + "-d", + "--config", + "-c", + "--verbose", + "-v", + ): assert flag in result.output, f"Missing flag {flag} in help" def test_model_required_without_list_flags(self, runner: CliRunner) -> None: @@ -189,6 +201,56 @@ def test_node_reduction_reported(self, runner: CliRunner, tmp_path: Path) -> Non assert result.exit_code == 0, result.output assert "10" in result.output assert "8" in result.output + assert "20.0% reduction" in result.output + + def test_node_increase_reported(self, runner: CliRunner, tmp_path: Path) -> None: + model_file = tmp_path / "model.onnx" + model_file.touch() + + original = _make_mock_model(num_nodes=10) + optimized = _make_mock_model(num_nodes=12) + with ( + patch(_LOAD_ONNX, return_value=original), + patch(_SAVE_ONNX), + patch(_OPTIMIZER) as mock_opt_cls, + ): + mock_opt_cls.return_value.optimize.return_value = optimized + result = runner.invoke(optimize, ["-m", str(model_file)]) + + assert result.exit_code == 0, result.output + assert "20.0% increase" in result.output + + def test_device_target_forwarded_to_optimizer( + self, runner: CliRunner, tmp_path: Path + ) -> None: + model_file = tmp_path / "model.onnx" + model_file.touch() + mock_model = _make_mock_model() + resolved_target = MagicMock(ep="DmlExecutionProvider", device="gpu") + resolved_ep_device = MagicMock() + + with ( + patch(_LOAD_ONNX, return_value=mock_model), + patch(_SAVE_ONNX), + patch(_OPTIMIZER) as mock_opt_cls, + patch( + "winml.modelkit.session.resolve_device", + return_value=resolved_target, + ) as mock_resolve, + patch("winml.modelkit.session.WinMLEPRegistry") as mock_registry_cls, + ): + mock_registry_cls.instance.return_value.auto_device.return_value = resolved_ep_device + mock_opt_cls.return_value.optimize.return_value = mock_model + result = runner.invoke(optimize, ["-m", str(model_file), "--device", "gpu"]) + + assert result.exit_code == 0, result.output + requested_target = mock_resolve.call_args.args[0] + assert requested_target.ep == "auto" + assert requested_target.device == "gpu" + assert ( + mock_opt_cls.return_value.optimize.call_args.kwargs["ep_device"] + is resolved_ep_device + ) # ============================================================================= diff --git a/tests/unit/optim/pipes/test_pipe_config.py b/tests/unit/optim/pipes/test_pipe_config.py index 4e6093daf..3e46b6bc0 100644 --- a/tests/unit/optim/pipes/test_pipe_config.py +++ b/tests/unit/optim/pipes/test_pipe_config.py @@ -22,6 +22,7 @@ import logging from typing import TYPE_CHECKING +from unittest.mock import MagicMock, patch from winml.modelkit.optim.pipes import ( GRAPH_CAPABILITIES, @@ -131,6 +132,44 @@ def test_verbose_parameter(self) -> None: assert config_verbose.verbose is True assert config_quiet.verbose is False + def test_ep_device_parameter(self) -> None: + ep_device = MagicMock() + + config = ORTGraphPipeConfig(ep_device=ep_device) + + assert config.ep_device is ep_device + + def test_build_config_forwards_ep_device(self) -> None: + ep_device = MagicMock() + + config = ORTGraphPipe.build_config(ep_device=ep_device) + + assert config.ep_device is ep_device + + def test_process_binds_resolved_ep_device(self) -> None: + model = MagicMock() + optimized_model = MagicMock() + ep_device = MagicMock() + handle = ep_device.device.ort_handle + session_options = MagicMock() + config = ORTGraphPipeConfig(ep_device=ep_device) + + with ( + patch("onnxruntime.SessionOptions", return_value=session_options), + patch("onnxruntime.InferenceSession") as inference_session, + patch("winml.modelkit.optim.pipes.graph.save_onnx"), + patch( + "winml.modelkit.optim.pipes.graph.load_onnx", + return_value=optimized_model, + ), + patch("winml.modelkit.session.lookup_device_spec", return_value=None), + ): + result = ORTGraphPipe().process(model, config) + + assert result is optimized_model + session_options.add_provider_for_devices.assert_called_once_with([handle], {}) + assert "providers" not in inference_session.call_args.kwargs + def test_always_disabled_optimizers(self) -> None: """AttentionFusion and EmbedLayerNormFusion are ALWAYS disabled. From f060afed9e8dcc71ad86c728a81fcaeffc10aa9b Mon Sep 17 00:00:00 2001 From: hualxie Date: Thu, 30 Jul 2026 17:33:31 +0800 Subject: [PATCH 2/4] Make optimization checks target aware --- src/winml/modelkit/commands/optimize.py | 85 ++++++++++++++++-------- src/winml/modelkit/optim/analysis.py | 23 ++++++- src/winml/modelkit/utils/cli.py | 9 ++- tests/unit/commands/test_optimize_cli.py | 43 ++++++++++++ tests/unit/optim/test_analysis.py | 59 ++++++++++++++++ 5 files changed, 188 insertions(+), 31 deletions(-) diff --git a/src/winml/modelkit/commands/optimize.py b/src/winml/modelkit/commands/optimize.py index 297177200..8b1c01bd4 100644 --- a/src/winml/modelkit/commands/optimize.py +++ b/src/winml/modelkit/commands/optimize.py @@ -228,7 +228,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 @@ -238,6 +257,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 @@ -246,6 +267,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]" ) @@ -257,7 +280,7 @@ def _run_check_optim(model: Path, all_caps: dict[str, Any], verbose: bool) -> No "[dim](this can take a while on large models)[/dim]" ) with console.status("[bold]Analyzing...[/bold]", spinner="dots"): - findings = analyze_model(onnx_model, all_caps) + findings = analyze_model(onnx_model, all_caps, ep_device=ep_device) _render_check_optim(console, findings, verbose) @@ -294,6 +317,7 @@ def _run_check_optim(model: Path, all_caps: dict[str, Any], verbose: bool) -> No 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( @@ -477,9 +501,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 @@ -543,31 +595,8 @@ def optimize( if verbose: optimizer_kwargs["verbose"] = True - if ep is not None or device is not None: - from ..session import ( - DeviceNotFound, - EPDeviceTarget, - UnknownListingPick, - WinMLEPNotDiscovered, - WinMLEPRegistrationFailed, - WinMLEPRegistry, - resolve_device, - ) - - try: - target = resolve_device( - EPDeviceTarget(ep=ep or "auto", device=(device or "auto").lower()) - ) - ep_device = WinMLEPRegistry.instance().auto_device(target) - except ( - DeviceNotFound, - RuntimeError, - UnknownListingPick, - ValueError, - WinMLEPNotDiscovered, - WinMLEPRegistrationFailed, - ) as e: - raise click.UsageError(f"Could not resolve optimization target: {e}") from e + if ep_device is not None: + assert target is not None optimizer_kwargs["ep_device"] = ep_device console.print( f"[bold blue]Target:[/bold blue] {target.ep} on {target.device.upper()}" diff --git a/src/winml/modelkit/optim/analysis.py b/src/winml/modelkit/optim/analysis.py index 649a5d2d3..f0b3fcb11 100644 --- a/src/winml/modelkit/optim/analysis.py +++ b/src/winml/modelkit/optim/analysis.py @@ -283,6 +283,7 @@ def _run_pipe(pipe: Any, model: ModelProto, config: Any) -> ModelProto: def _iter_findings( model: ModelProto, capabilities: dict[str, CapabilityDef], + **optimizer_kwargs: Any, ) -> Iterator[tuple[CapabilityFinding, ModelProto]]: """Yield ``(finding, produced_model)`` for every applicable optimization. @@ -305,6 +306,7 @@ def _iter_findings( # 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: @@ -340,6 +342,14 @@ def _iter_findings( ] for cap_name, cap in probe_caps: + 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 = ep_device.device.ep_name + if not any(normalize_ep_name(name) == target_ep for name in cap.ep_constraint): + continue + # Enable only this capability (plus its dependencies) on top of the # all-defaults configuration. kebab = dict(kebab_defaults) @@ -350,6 +360,7 @@ def _iter_findings( 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) @@ -406,6 +417,7 @@ def _iter_findings( def analyze_model( model: ModelProto, capabilities: dict[str, CapabilityDef], + **optimizer_kwargs: Any, ) -> list[CapabilityFinding]: """Probe every applicable optimization capability against ``model``. @@ -418,17 +430,23 @@ 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``. Returns: Applicable findings in pipeline order, each naming the affected nodes and constants. """ - return [finding for finding, _ in _iter_findings(model, capabilities)] + return [ + finding + for finding, _ in _iter_findings(model, capabilities, **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. @@ -442,10 +460,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) diff --git a/src/winml/modelkit/utils/cli.py b/src/winml/modelkit/utils/cli.py index 46da14c8a..d1a12049c 100644 --- a/src/winml/modelkit/utils/cli.py +++ b/src/winml/modelkit/utils/cli.py @@ -415,6 +415,7 @@ def ep_option( default: str | None = None, include_auto: bool = False, include_all: bool = False, + include_cuda: bool = False, ) -> Callable[[F], F]: """Add --ep (execution provider) option to a Click command. @@ -428,6 +429,8 @@ def ep_option( (default: False). include_all: Whether to include "all" as a valid choice (default: False). + include_cuda: Whether to include CUDA aliases and the full provider name + (default: False). Returns: Decorator function @@ -440,7 +443,11 @@ def ep_option( if optional_message: help_text = f"{help_text}. {optional_message}" - ep_choices = [name for name in ALL_EP_NAMES if name not in ("cuda", "CUDAExecutionProvider")] + ep_choices = [ + name + for name in ALL_EP_NAMES + if include_cuda or name not in ("cuda", "CUDAExecutionProvider") + ] choices = ["auto", *ep_choices] if include_auto else ep_choices choices = ["all", *choices] if include_all else choices diff --git a/tests/unit/commands/test_optimize_cli.py b/tests/unit/commands/test_optimize_cli.py index 58586c1a2..feef461ef 100644 --- a/tests/unit/commands/test_optimize_cli.py +++ b/tests/unit/commands/test_optimize_cli.py @@ -332,6 +332,49 @@ def test_check_optim_no_findings_message(self, runner: CliRunner, tmp_path: Path assert result.exit_code == 0, result.output assert "No registered optimizations" in result.output + def test_check_optim_forwards_resolved_target( + self, runner: CliRunner, tmp_path: Path + ) -> None: + model_file = tmp_path / "model.onnx" + model_file.touch() + target = MagicMock(ep="DmlExecutionProvider", device="gpu") + ep_device = MagicMock() + + with ( + patch(_LOAD_ONNX, return_value=_make_mock_model()), + patch(_ANALYZE_MODEL, return_value=[]) as mock_analyze, + patch( + "winml.modelkit.commands.optimize._resolve_optimization_target", + return_value=(target, ep_device), + ) as mock_resolve, + ): + result = runner.invoke( + optimize, + ["-m", str(model_file), "--check-optim", "--device", "gpu"], + ) + + assert result.exit_code == 0, result.output + mock_resolve.assert_called_once_with(None, "gpu") + assert mock_analyze.call_args.kwargs["ep_device"] is ep_device + assert "DmlExecutionProvider on GPU" in result.output + + def test_check_optim_accepts_cuda_ep(self, runner: CliRunner, tmp_path: Path) -> None: + model_file = tmp_path / "model.onnx" + model_file.touch() + + with patch( + "winml.modelkit.commands.optimize._resolve_optimization_target", + side_effect=ValueError("CUDA unavailable"), + ) as mock_resolve: + result = runner.invoke( + optimize, + ["-m", str(model_file), "--check-optim", "--ep", "cuda"], + ) + + assert result.exit_code != 0 + mock_resolve.assert_called_once_with("cuda", None) + assert "Invalid value for '--ep'" not in result.output + class TestConfigFile: """Config file loading and precedence.""" diff --git a/tests/unit/optim/test_analysis.py b/tests/unit/optim/test_analysis.py index a69856d17..1c7c59190 100644 --- a/tests/unit/optim/test_analysis.py +++ b/tests/unit/optim/test_analysis.py @@ -12,6 +12,8 @@ from __future__ import annotations +from unittest.mock import MagicMock + import numpy as np from onnx import GraphProto, ModelProto, TensorProto, helper, numpy_helper @@ -29,6 +31,7 @@ _diff_nodes, ) from winml.modelkit.optim.pipes import get_all_capabilities +from winml.modelkit.optim.registry import BoolCapability, CapabilityCategory # ============================================================================= @@ -403,6 +406,62 @@ def test_input_model_not_mutated(self) -> None: after_value = numpy_helper.to_array(model.graph.initializer[0]) np.testing.assert_array_equal(before_value, after_value) + def test_target_context_forwarded_and_ep_constraints_filtered( + self, monkeypatch + ) -> None: + dml_cap = BoolCapability( + name="dml-only", + ort_name=None, + description="DML-only probe", + category=CapabilityCategory.MISC, + ep_constraint=("DML",), + ) + cuda_cap = BoolCapability( + name="cuda-only", + ort_name=None, + description="CUDA-only probe", + category=CapabilityCategory.MISC, + ep_constraint=("CUDA",), + ) + cap_registry = {cap.name: cap for cap in (dml_cap, cuda_cap)} + captured_ep_devices = [] + + class TargetAwarePipe: + name = "target-aware" + capabilities = cap_registry + + @classmethod + def build_config(cls, **kwargs): + captured_ep_devices.append(kwargs.get("ep_device")) + return kwargs + + @staticmethod + def process(model, config): + enabled = next( + (name for name in ("dml_only", "cuda_only") if config.get(name)), + None, + ) + if enabled is not None: + model.graph.node.append( + helper.make_node("Identity", ["z"], [f"{enabled}_output"]) + ) + return model + + ep_device = MagicMock() + ep_device.device.ep_name = "DmlExecutionProvider" + monkeypatch.setattr("winml.modelkit.optim.pipes.PIPES", [TargetAwarePipe]) + monkeypatch.setattr("winml.modelkit.onnx.infer_shapes", lambda model: model) + + findings = analyze_model( + _benign_model(), + cap_registry, + ep_device=ep_device, + ) + + assert [finding.name for finding in findings] == ["dml-only"] + assert captured_ep_devices + assert all(value is ep_device for value in captured_ep_devices) + # ============================================================================= # OPTIMIZATION OUTPUT ITERATION From 9656141df14eff06c67965259db0653c4cd009be Mon Sep 17 00:00:00 2001 From: hualxie Date: Fri, 31 Jul 2026 10:47:34 +0800 Subject: [PATCH 3/4] Add VAE graph optimization surgeries Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../modelkit/optim/capabilities/surgery.py | 44 +- src/winml/modelkit/optim/pipes/__init__.py | 20 +- src/winml/modelkit/optim/pipes/surgery.py | 1053 ++++++++++++++++- tests/unit/optim/pipes/test_pipe_surgery.py | 6 + .../test_pipe_surgery_graph_reductions.py | 624 ++++++++++ tests/unit/optim/test_optimizer.py | 7 +- vae_opt.md | 484 ++++++++ 7 files changed, 2216 insertions(+), 22 deletions(-) create mode 100644 tests/unit/optim/pipes/test_pipe_surgery_graph_reductions.py create mode 100644 vae_opt.md diff --git a/src/winml/modelkit/optim/capabilities/surgery.py b/src/winml/modelkit/optim/capabilities/surgery.py index 0b6ec0768..97c8d28e7 100644 --- a/src/winml/modelkit/optim/capabilities/surgery.py +++ b/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",), +) diff --git a/src/winml/modelkit/optim/pipes/__init__.py b/src/winml/modelkit/optim/pipes/__init__.py index 9315b71fa..c30f69b6d 100644 --- a/src/winml/modelkit/optim/pipes/__init__.py +++ b/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,13 +19,19 @@ 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 @@ -34,6 +39,7 @@ # - 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", diff --git a/src/winml/modelkit/optim/pipes/surgery.py b/src/winml/modelkit/optim/pipes/surgery.py index 90a95f46a..75c7669e0 100644 --- a/src/winml/modelkit/optim/pipes/surgery.py +++ b/src/winml/modelkit/optim/pipes/surgery.py @@ -5,8 +5,8 @@ """Surgery pipe for precise model modifications. This pipe performs targeted graph transformations that are not part of -ONNX Runtime's standard optimization passes. Surgery operations run before -ORT optimizations to prepare models for quantization or specific execution providers. +ONNX Runtime's standard optimization passes. Surgery operations run before or +after ORT optimization according to the graph evidence they require. Use cases: - Clamp extreme constant values to prevent quantization issues @@ -17,17 +17,15 @@ import logging from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, ClassVar +from typing import Any, ClassVar import numpy as np +import onnx from ..capabilities import surgery from .base import BasePipe, PipeConfig, caps_dict -if TYPE_CHECKING: - import onnx - logger = logging.getLogger(__name__) @@ -35,12 +33,24 @@ # MODULE-LEVEL CAPABILITIES # ============================================================================= -SURGERY_CAPABILITIES: dict[str, Any] = caps_dict( +PRE_SURGERY_CAPABILITIES: dict[str, Any] = caps_dict( + surgery.SIMPLIFY_L2_NORMALIZATION, + surgery.GATHERND_TO_RESIZE, + surgery.SCALED_MATMUL_TO_FUSED_MATMUL, +) + +POST_SURGERY_CAPABILITIES: dict[str, Any] = caps_dict( surgery.CLAMP_CONSTANT_VALUES, surgery.REMOVE_ISNAN_IN_ATTENTION_MASK, surgery.UNTIE_CONSTANT_BATCHED_MATMUL, + surgery.SILU_TO_QUICK_GELU, ) +SURGERY_CAPABILITIES: dict[str, Any] = { + **PRE_SURGERY_CAPABILITIES, + **POST_SURGERY_CAPABILITIES, +} + # ============================================================================= # SURGERYPIPECONFIG @@ -60,6 +70,13 @@ class SurgeryPipeConfig(PipeConfig): mask_value: Replacement value for -inf (default: -1e3) untie_constant_batched_matmul: Make a batched MatMul's constant operand runtime-valued so OpenVINO GPU can select a gemm implementation + simplify_l2_normalization: Remove redundant Clip and Expand nodes from + exact ReduceL2-based normalization patterns + gathernd_to_resize: Replace exact 2x nearest-neighbor GatherND + upsampling grids with standard ONNX Resize + silu_to_quick_gelu: Replace x*Sigmoid(x) with the Microsoft QuickGelu op + scaled_matmul_to_fused_matmul: Move a scalar Mul into Microsoft + FusedMatMul's alpha attribute verbose: Enable verbose logging """ @@ -68,9 +85,631 @@ class SurgeryPipeConfig(PipeConfig): clamp_max: float = 1e3 remove_isnan_in_attention_mask: bool = False untie_constant_batched_matmul: bool = False + simplify_l2_normalization: bool = False + gathernd_to_resize: bool = False + silu_to_quick_gelu: bool = False + scaled_matmul_to_fused_matmul: bool = False verbose: bool = False +@dataclass +class _GraphIndex: + """Producer, consumer, constant, type, and output metadata for one graph.""" + + producers: dict[str, onnx.NodeProto] + consumers: dict[str, list[onnx.NodeProto]] + initializers: dict[str, onnx.TensorProto] + element_types: dict[str, int] + shapes: dict[str, tuple[int | str | None, ...]] + graph_outputs: set[str] + + @classmethod + def build(cls, model: onnx.ModelProto) -> _GraphIndex: + graph = model.graph + producers: dict[str, onnx.NodeProto] = {} + consumers: dict[str, list[onnx.NodeProto]] = {} + for node in graph.node: + for output in node.output: + if output: + producers[output] = node + consumed_names = {name for name in node.input if name} + for attribute in node.attribute: + if attribute.type == onnx.AttributeProto.GRAPH: + consumed_names.update(_captured_tensor_names(attribute.g)) + elif attribute.type == onnx.AttributeProto.GRAPHS: + for nested_graph in attribute.graphs: + consumed_names.update(_captured_tensor_names(nested_graph)) + for name in consumed_names: + consumers.setdefault(name, []).append(node) + + initializers = {initializer.name: initializer for initializer in graph.initializer} + element_types: dict[str, int] = { + initializer.name: int(initializer.data_type) for initializer in graph.initializer + } + for value in (*graph.input, *graph.value_info, *graph.output): + if value.type.HasField("tensor_type"): + element_types[value.name] = int(value.type.tensor_type.elem_type) + shapes: dict[str, tuple[int | str | None, ...]] = {} + for value in (*graph.input, *graph.value_info, *graph.output): + if not value.type.HasField("tensor_type"): + continue + tensor_type = value.type.tensor_type + if not tensor_type.HasField("shape"): + continue + dimensions: list[int | str | None] = [] + for dimension in tensor_type.shape.dim: + if dimension.HasField("dim_value"): + dimensions.append(int(dimension.dim_value)) + elif dimension.HasField("dim_param") and dimension.dim_param: + dimensions.append(dimension.dim_param) + else: + dimensions.append(None) + shapes[value.name] = tuple(dimensions) + for initializer in graph.initializer: + shapes.setdefault(initializer.name, tuple(int(dim) for dim in initializer.dims)) + + return cls( + producers=producers, + consumers=consumers, + initializers=initializers, + element_types=element_types, + shapes=shapes, + graph_outputs={output.name for output in graph.output if output.name}, + ) + + +class _NameAllocator: + """Allocate graph-wide unique tensor and node names.""" + + def __init__(self, model: onnx.ModelProto) -> None: + graph = model.graph + self._used = { + name + for name in ( + [initializer.name for initializer in graph.initializer] + + [value.name for value in (*graph.input, *graph.value_info, *graph.output)] + + [node.name for node in graph.node] + + [name for node in graph.node for name in (*node.input, *node.output)] + ) + if name + } + + def new(self, prefix: str) -> str: + candidate = prefix + suffix = 0 + while candidate in self._used: + suffix += 1 + candidate = f"{prefix}_{suffix}" + self._used.add(candidate) + return candidate + + +def _captured_tensor_names(graph: onnx.GraphProto) -> set[str]: + """Return outer-scope tensor names referenced by a nested graph.""" + locally_defined = {value.name for value in graph.input if value.name} + locally_defined.update(initializer.name for initializer in graph.initializer) + locally_defined.update(output for node in graph.node for output in node.output if output) + referenced = {name for node in graph.node for name in node.input if name} + for node in graph.node: + for attribute in node.attribute: + if attribute.type == onnx.AttributeProto.GRAPH: + referenced.update(_captured_tensor_names(attribute.g)) + elif attribute.type == onnx.AttributeProto.GRAPHS: + for nested_graph in attribute.graphs: + referenced.update(_captured_tensor_names(nested_graph)) + return referenced - locally_defined + + +def _referenced_tensor_names(graph: onnx.GraphProto) -> set[str]: + """Collect tensor names referenced by a graph and its nested graphs.""" + referenced = {value.name for value in (*graph.input, *graph.output) if value.name} + for node in graph.node: + referenced.update(name for name in node.input if name) + for attribute in node.attribute: + if attribute.type == onnx.AttributeProto.GRAPH: + referenced.update(_referenced_tensor_names(attribute.g)) + elif attribute.type == onnx.AttributeProto.GRAPHS: + for nested_graph in attribute.graphs: + referenced.update(_referenced_tensor_names(nested_graph)) + return referenced + + +def _attribute(node: onnx.NodeProto, name: str, default: Any = None) -> Any: + for attribute in node.attribute: + if attribute.name == name: + return onnx.helper.get_attribute_value(attribute) + return default + + +def _tensor_array(tensor: onnx.TensorProto) -> np.ndarray | None: + try: + values = np.asarray(onnx.numpy_helper.to_array(tensor)) + except (TypeError, ValueError): + return None + if tensor.data_type == onnx.TensorProto.BFLOAT16 and np.issubdtype(values.dtype, np.integer): + values = np.left_shift(values.astype(np.uint32), 16).view(np.float32) + return values + + +def _constant_array(index: _GraphIndex, name: str) -> np.ndarray | None: + """Read an initializer or tensor-valued Constant as a NumPy array.""" + initializer = index.initializers.get(name) + if initializer is not None: + return _tensor_array(initializer) + + producer = index.producers.get(name) + if producer is None or producer.op_type != "Constant" or not _is_default_domain(producer): + return None + value = _attribute(producer, "value") + if value is not None: + return _tensor_array(value) + for attribute_name in ("value_float", "value_floats", "value_int", "value_ints"): + attribute_value = _attribute(producer, attribute_name) + if attribute_value is not None: + return np.asarray(attribute_value) + return None + + +def _constant_scalar(index: _GraphIndex, name: str) -> float | int | None: + values = _constant_array(index, name) + if values is None or values.size != 1: + return None + return values.reshape(-1)[0].item() + + +def _constant_element_type(index: _GraphIndex, name: str) -> int | None: + initializer = index.initializers.get(name) + if initializer is not None: + return int(initializer.data_type) + producer = index.producers.get(name) + if producer is None or producer.op_type != "Constant" or not _is_default_domain(producer): + return index.element_types.get(name) + value = _attribute(producer, "value") + if isinstance(value, onnx.TensorProto): + return int(value.data_type) + if ( + _attribute(producer, "value_float") is not None + or _attribute(producer, "value_floats") is not None + ): + return int(onnx.TensorProto.FLOAT) + if ( + _attribute(producer, "value_int") is not None + or _attribute(producer, "value_ints") is not None + ): + return int(onnx.TensorProto.INT64) + return index.element_types.get(name) + + +def _is_default_domain(node: onnx.NodeProto) -> bool: + return node.domain in ("", "ai.onnx") + + +def _node_output(node: onnx.NodeProto) -> str | None: + return node.output[0] if len(node.output) == 1 and node.output[0] else None + + +def _only_consumer( + index: _GraphIndex, + tensor_name: str, + op_type: str, +) -> onnx.NodeProto | None: + consumers = index.consumers.get(tensor_name, []) + if len(consumers) != 1: + return None + consumer = consumers[0] + if consumer.op_type != op_type or not _is_default_domain(consumer): + return None + return consumer + + +def _constant_axes(index: _GraphIndex, node: onnx.NodeProto) -> list[int] | None: + axes = _attribute(node, "axes") + if axes is None and len(node.input) > 1 and node.input[1]: + values = _constant_array(index, node.input[1]) + if values is None or not np.issubdtype(values.dtype, np.integer): + return None + axes = values.reshape(-1).tolist() + if axes is None: + return None + return [int(axis) for axis in np.asarray(axes).reshape(-1).tolist()] + + +def _remove_nodes(model: onnx.ModelProto, remove_ids: set[int]) -> None: + remaining = [node for node in model.graph.node if id(node) not in remove_ids] + del model.graph.node[:] + model.graph.node.extend(remaining) + + +def _prune_dead_ancestors(model: onnx.ModelProto, seed_names: set[str]) -> None: + """Remove only dead producers and constants reached from removed-node inputs.""" + pending = {name for name in seed_names if name} + candidate_initializers: set[str] = set() + while pending: + index = _GraphIndex.build(model) + remove_ids: set[int] = set() + next_pending: set[str] = set() + for name in pending: + if name in index.initializers: + candidate_initializers.add(name) + producer = index.producers.get(name) + if producer is None: + continue + if any( + output in index.graph_outputs or index.consumers.get(output) + for output in producer.output + if output + ): + continue + remove_ids.add(id(producer)) + next_pending.update(input_name for input_name in producer.input if input_name) + if not remove_ids: + break + _remove_nodes(model, remove_ids) + pending = next_pending + + referenced = _referenced_tensor_names(model.graph) + remaining_initializers = [ + initializer + for initializer in model.graph.initializer + if initializer.name not in candidate_initializers or initializer.name in referenced + ] + del model.graph.initializer[:] + model.graph.initializer.extend(remaining_initializers) + + +def _prune_stale_value_info(model: onnx.ModelProto) -> None: + graph = model.graph + defined = {value.name for value in (*graph.input, *graph.output) if value.name} + defined.update(initializer.name for initializer in graph.initializer) + defined.update(output for node in graph.node for output in node.output if output) + remaining = [value for value in graph.value_info if value.name in defined] + del graph.value_info[:] + graph.value_info.extend(remaining) + + +def _ensure_ms_opset(model: onnx.ModelProto) -> None: + for opset in model.opset_import: + if opset.domain == "com.microsoft": + if opset.version < 1: + opset.version = 1 + return + model.opset_import.append(onnx.helper.make_opsetid("com.microsoft", 1)) + + +def _default_opset(model: onnx.ModelProto) -> int: + return next( + (int(opset.version) for opset in model.opset_import if opset.domain in ("", "ai.onnx")), + 0, + ) + + +def _other_input_with_constant( + index: _GraphIndex, + node: onnx.NodeProto, + expected: float, +) -> str | None: + if len(node.input) != 2: + return None + for constant_index in (0, 1): + value = _constant_scalar(index, node.input[constant_index]) + if value is not None and float(value) == expected: + return node.input[1 - constant_index] + return None + + +def _expand_preserves_shape( + index: _GraphIndex, + expand: onnx.NodeProto, + source_name: str, +) -> bool: + """Prove that Expand's target shape is exactly the source tensor shape.""" + if len(expand.input) != 2: + return False + target_name = expand.input[1] + target_producer = _producer(index, target_name, "Shape") + if ( + target_producer is not None + and target_producer.input + and target_producer.input[0] == source_name + and int(_attribute(target_producer, "start", 0)) == 0 + and _attribute(target_producer, "end") is None + ): + return True + + source_shape = index.shapes.get(source_name) + expanded_output = _node_output(expand) + expanded_shape = index.shapes.get(expanded_output) if expanded_output else None + if ( + source_shape is not None + and expanded_shape is not None + and None not in source_shape + and source_shape == expanded_shape + ): + return True + + target_shape = _constant_array(index, target_name) + return ( + source_shape is not None + and None not in source_shape + and target_shape is not None + and np.issubdtype(target_shape.dtype, np.integer) + and tuple(int(value) for value in target_shape.reshape(-1)) == source_shape + ) + + +def _producer( + index: _GraphIndex, + tensor_name: str, + op_type: str, +) -> onnx.NodeProto | None: + node = index.producers.get(tensor_name) + if node is None or node.op_type != op_type or not _is_default_domain(node): + return None + return node + + +def _match_dimension_scalar(index: _GraphIndex, name: str) -> str | None: + """Return the one-element shape vector representing a dimension scalar.""" + squeeze = _producer(index, name, "Squeeze") + if squeeze is not None and squeeze.input: + axes = _constant_axes(index, squeeze) + if axes is None or axes in ([0], [-1]): + shape_vector = squeeze.input[0] + shape = _producer(index, shape_vector, "Shape") + if shape is not None: + start = int(_attribute(shape, "start", 0)) + end = _attribute(shape, "end") + if end is not None and int(end) == start + 1: + return shape_vector + + reshaped_vectors = [] + for consumer in index.consumers.get(name, []): + if ( + consumer.op_type != "Reshape" + or not _is_default_domain(consumer) + or len(consumer.input) < 2 + or consumer.input[0] != name + ): + continue + target_shape = _constant_array(index, consumer.input[1]) + output = _node_output(consumer) + if ( + target_shape is not None + and target_shape.size == 1 + and int(target_shape.reshape(-1)[0]) == -1 + and output is not None + ): + reshaped_vectors.append(output) + return reshaped_vectors[0] if len(reshaped_vectors) == 1 else None + + +def _match_nearest_index_vector(index: _GraphIndex, name: str) -> str | None: + """Match floor((arange(2*d) + 0.5) * (d / (2*d))).""" + cast_int = _producer(index, name, "Cast") + if cast_int is None or int(_attribute(cast_int, "to", -1)) != int(onnx.TensorProto.INT64): + return None + scaled = _producer(index, cast_int.input[0], "Mul") if cast_int.input else None + if scaled is None or len(scaled.input) != 2: + return None + + left = index.producers.get(scaled.input[0]) + right = index.producers.get(scaled.input[1]) + if ( + left is not None + and right is not None + and left.op_type == "Add" + and right.op_type == "Div" + and _is_default_domain(left) + and _is_default_domain(right) + ): + offset, ratio = left, right + elif ( + left is not None + and right is not None + and left.op_type == "Div" + and right.op_type == "Add" + and _is_default_domain(left) + and _is_default_domain(right) + ): + offset, ratio = right, left + else: + return None + + range_name = _other_input_with_constant(index, offset, 0.5) + arange = _producer(index, range_name, "Range") if range_name else None + if arange is None or len(arange.input) != 3: + return None + if _constant_scalar(index, arange.input[0]) != 0: + return None + if _constant_scalar(index, arange.input[2]) != 1: + return None + + end_cast = _producer(index, arange.input[1], "Cast") + if end_cast is None or int(_attribute(end_cast, "to", -1)) != int(onnx.TensorProto.FLOAT): + return None + doubled_size = _producer(index, end_cast.input[0], "Mul") if end_cast.input else None + if doubled_size is None: + return None + dimension_scalar = _other_input_with_constant(index, doubled_size, 2.0) + if dimension_scalar is None: + return None + shape_vector = _match_dimension_scalar(index, dimension_scalar) + if shape_vector is None: + return None + + if len(ratio.input) != 2: + return None + dimension_float = _producer(index, ratio.input[0], "Cast") + if ( + dimension_float is None + or int(_attribute(dimension_float, "to", -1)) != int(onnx.TensorProto.FLOAT) + or not dimension_float.input + or dimension_float.input[0] != dimension_scalar + ): + return None + denominator = _producer(index, ratio.input[1], "Mul") + if denominator is None: + return None + denominator_other = _other_input_with_constant(index, denominator, 2.0) + if denominator_other != dimension_float.output[0]: + return None + return shape_vector + + +def _match_unsqueeze( + index: _GraphIndex, + tensor_name: str, + axes: list[int], +) -> onnx.NodeProto | None: + node = _producer(index, tensor_name, "Unsqueeze") + if node is None or _constant_axes(index, node) != axes: + return None + return node + + +def _match_resize_indices(index: _GraphIndex, name: str) -> tuple[str, str] | None: + """Match a broadcasted pair of exact 2x nearest-neighbor index vectors.""" + concat = _producer(index, name, "Concat") + if concat is None or len(concat.input) != 2 or int(_attribute(concat, "axis", 0)) != -1: + return None + + outer_h = _match_unsqueeze(index, concat.input[0], [-1]) + outer_w = _match_unsqueeze(index, concat.input[1], [-1]) + if outer_h is None or outer_w is None: + return None + expand_h = _producer(index, outer_h.input[0], "Expand") + expand_w = _producer(index, outer_w.input[0], "Expand") + if ( + expand_h is None + or expand_w is None + or len(expand_h.input) != 2 + or len(expand_w.input) != 2 + or expand_h.input[1] != expand_w.input[1] + ): + return None + + inner_h = _match_unsqueeze(index, expand_h.input[0], [-1]) + if inner_h is None or not inner_h.input: + return None + h_indices = inner_h.input[0] + w_indices = expand_w.input[0] + h_shape_vector = _match_nearest_index_vector(index, h_indices) + w_shape_vector = _match_nearest_index_vector(index, w_indices) + if h_shape_vector is None or w_shape_vector is None: + return None + + grid_shape = _producer(index, expand_h.input[1], "Shape") + grid_max = ( + _producer(index, grid_shape.input[0], "Max") + if grid_shape is not None and grid_shape.input + else None + ) + if grid_max is None or len(grid_max.input) != 2: + return None + if set(grid_max.input) != {inner_h.output[0], w_indices}: + return None + return h_shape_vector, w_shape_vector + + +@dataclass(frozen=True) +class _ResizeMatch: + pre_cast: onnx.NodeProto + pre_transpose: onnx.NodeProto + gather: onnx.NodeProto + post_transpose: onnx.NodeProto + post_cast: onnx.NodeProto + source: str + indices: str + + +def _match_resize_path(index: _GraphIndex, gather: onnx.NodeProto) -> _ResizeMatch | None: + """Match the exact FP16/FP32 GatherND path used for 2x NCHW upsampling.""" + if ( + gather.op_type != "GatherND" + or not _is_default_domain(gather) + or len(gather.input) != 2 + or len(gather.output) != 1 + or int(_attribute(gather, "batch_dims", 0)) != 0 + ): + return None + + pre_transpose = _producer(index, gather.input[0], "Transpose") + if pre_transpose is None or len(pre_transpose.input) != 1: + return None + pre_cast = _producer(index, pre_transpose.input[0], "Cast") + if ( + pre_cast is None + or len(pre_cast.input) != 1 + or int(_attribute(pre_cast, "to", -1)) != int(onnx.TensorProto.FLOAT) + ): + return None + + post_transpose = _only_consumer(index, gather.output[0], "Transpose") + if post_transpose is None: + return None + post_output = _node_output(post_transpose) + post_cast = _only_consumer(index, post_output, "Cast") if post_output else None + if post_cast is None or len(post_cast.output) != 1: + return None + + permutation = [2, 3, 0, 1] + if list(_attribute(pre_transpose, "perm", [])) != permutation: + return None + if list(_attribute(post_transpose, "perm", [])) != permutation: + return None + + pre_cast_output = _node_output(pre_cast) + pre_transpose_output = _node_output(pre_transpose) + if ( + pre_cast_output is None + or _only_consumer(index, pre_cast_output, "Transpose") is not pre_transpose + or pre_transpose_output is None + or _only_consumer(index, pre_transpose_output, "GatherND") is not gather + ): + return None + + source = pre_cast.input[0] + source_type = index.element_types.get(source) + if source_type not in (int(onnx.TensorProto.FLOAT), int(onnx.TensorProto.FLOAT16)): + return None + if int(_attribute(post_cast, "to", -1)) != source_type: + return None + + intermediate_outputs = { + pre_cast.output[0], + pre_transpose.output[0], + gather.output[0], + post_transpose.output[0], + } + if intermediate_outputs & index.graph_outputs: + return None + + shape_vectors = _match_resize_indices(index, gather.input[1]) + if shape_vectors is None: + return None + reshape = _producer(index, source, "Reshape") + if reshape is None or len(reshape.input) < 2: + return None + shape_concat = _producer(index, reshape.input[1], "Concat") + if ( + shape_concat is None + or len(shape_concat.input) != 4 + or int(_attribute(shape_concat, "axis", 0)) != 0 + or tuple(shape_concat.input[2:4]) != shape_vectors + ): + return None + + return _ResizeMatch( + pre_cast=pre_cast, + pre_transpose=pre_transpose, + gather=gather, + post_transpose=post_transpose, + post_cast=post_cast, + source=source, + indices=gather.input[1], + ) + + # ============================================================================= # SURGERYPIPE # ============================================================================= @@ -79,16 +718,14 @@ class SurgeryPipeConfig(PipeConfig): class SurgeryPipe(BasePipe[SurgeryPipeConfig]): """Surgery pipe for precise model modifications. - This pipe performs targeted graph transformations to prepare models - for quantization or specific execution providers. It runs before - ORT optimizations. + This pipe performs targeted post-optimization graph transformations. Currently supported operations: - clamp-constant-values: Clamp extreme float constants (e.g., -inf → -1e3) """ name: ClassVar[str] = "surgery" - capabilities: ClassVar[dict[str, Any]] = SURGERY_CAPABILITIES + capabilities: ClassVar[dict[str, Any]] = POST_SURGERY_CAPABILITIES @classmethod def build_config(cls, **kwargs: Any) -> SurgeryPipeConfig: @@ -100,6 +737,10 @@ def build_config(cls, **kwargs: Any) -> SurgeryPipeConfig: - clamp_min: Minimum value for clamping (default: -1e3) - clamp_max: Maximum value for clamping (default: 1e3) - remove_isnan_in_attention_mask: Remove IsNaN guard patterns + - simplify_l2_normalization: Remove redundant normalization ops + - gathernd_to_resize: Replace exact GatherND upsampling grids + - silu_to_quick_gelu: Emit QuickGelu(alpha=1) for exact SiLU + - scaled_matmul_to_fused_matmul: Emit FusedMatMul for scaled MatMul - verbose: Enable verbose logging Returns: @@ -111,6 +752,10 @@ def build_config(cls, **kwargs: Any) -> SurgeryPipeConfig: clamp_max=kwargs.get("clamp_max", 1e3), remove_isnan_in_attention_mask=kwargs.get("remove_isnan_in_attention_mask", False), untie_constant_batched_matmul=kwargs.get("untie_constant_batched_matmul", False), + simplify_l2_normalization=kwargs.get("simplify_l2_normalization", False), + gathernd_to_resize=kwargs.get("gathernd_to_resize", False), + silu_to_quick_gelu=kwargs.get("silu_to_quick_gelu", False), + scaled_matmul_to_fused_matmul=kwargs.get("scaled_matmul_to_fused_matmul", False), verbose=kwargs.get("verbose", False), ) @@ -128,6 +773,10 @@ def should_process(cls, config: SurgeryPipeConfig) -> bool: config.clamp_constant_values or config.remove_isnan_in_attention_mask or config.untie_constant_batched_matmul + or config.simplify_l2_normalization + or config.gathernd_to_resize + or config.silu_to_quick_gelu + or config.scaled_matmul_to_fused_matmul ) def process(self, model: onnx.ModelProto, config: SurgeryPipeConfig) -> onnx.ModelProto: @@ -142,7 +791,6 @@ def process(self, model: onnx.ModelProto, config: SurgeryPipeConfig) -> onnx.Mod """ if not self.should_process(config): return model - # Import onnx inside method to avoid import errors import onnx @@ -161,6 +809,18 @@ def process(self, model: onnx.ModelProto, config: SurgeryPipeConfig) -> onnx.Mod if config.untie_constant_batched_matmul: model_copy = self._untie_constant_batched_matmul(model_copy, config.verbose) + if config.simplify_l2_normalization: + model_copy = self._simplify_l2_normalization(model_copy, config.verbose) + + if config.gathernd_to_resize: + model_copy = self._gathernd_to_resize(model_copy, config.verbose) + + if config.silu_to_quick_gelu: + model_copy = self._silu_to_quick_gelu(model_copy, config.verbose) + + if config.scaled_matmul_to_fused_matmul: + model_copy = self._scaled_matmul_to_fused_matmul(model_copy, config.verbose) + return model_copy def _clamp_constant_values( @@ -476,3 +1136,372 @@ def zero_for(dtype: int) -> str: ) return model + + # ----------------------------------------------------------------- + # simplify-l2-normalization + # ----------------------------------------------------------------- + + def _simplify_l2_normalization( + self, + model: onnx.ModelProto, + verbose: bool = False, + ) -> onnx.ModelProto: + """Remove Clip(min=0) and Expand from exact ReduceL2 normalization.""" + index = _GraphIndex.build(model) + remove_ids: set[int] = set() + prune_seeds: set[str] = set() + simplified = 0 + + for reduce_l2 in list(model.graph.node): + if ( + reduce_l2.op_type != "ReduceL2" + or not _is_default_domain(reduce_l2) + or not reduce_l2.input + or _constant_axes(index, reduce_l2) != [1] + or int(_attribute(reduce_l2, "keepdims", 1)) != 1 + ): + continue + norm_output = _node_output(reduce_l2) + if norm_output is None: + continue + clip = _only_consumer(index, norm_output, "Clip") + if clip is None or _node_output(clip) in index.graph_outputs: + continue + + clip_min_values: list[float] = [] + attribute_min = _attribute(clip, "min") + if attribute_min is not None: + clip_min_values.append(float(attribute_min)) + if len(clip.input) > 1 and clip.input[1]: + input_min = _constant_scalar(index, clip.input[1]) + if input_min is None: + continue + clip_min_values.append(float(input_min)) + if not clip_min_values or any(value != 0.0 for value in clip_min_values): + continue + if _attribute(clip, "max") is not None: + continue + if len(clip.input) > 2 and clip.input[2]: + continue + + clip_output = _node_output(clip) + expand = _only_consumer(index, clip_output, "Expand") if clip_output else None + if ( + expand is None + or len(expand.input) != 2 + or _node_output(expand) in index.graph_outputs + or not _expand_preserves_shape(index, expand, reduce_l2.input[0]) + ): + continue + expand_output = _node_output(expand) + div = _only_consumer(index, expand_output, "Div") if expand_output else None + if ( + div is None + or len(div.input) != 2 + or div.input[0] != reduce_l2.input[0] + or div.input[1] != expand_output + ): + continue + + div.input[1] = norm_output + remove_ids.update((id(clip), id(expand))) + prune_seeds.update(name for node in (clip, expand) for name in node.input if name) + simplified += 1 + if verbose: + logger.info( + " simplify-l2-normalization: %s -> %s now broadcasts directly", + reduce_l2.name or norm_output, + div.name or div.output[0], + ) + + if not simplified: + return model + + _remove_nodes(model, remove_ids) + _prune_dead_ancestors(model, prune_seeds) + _prune_stale_value_info(model) + logger.info( + "SurgeryPipe: simplify-l2-normalization: simplified %d normalization(s)", + simplified, + ) + return model + + # ----------------------------------------------------------------- + # gathernd-to-resize + # ----------------------------------------------------------------- + + def _gathernd_to_resize( + self, + model: onnx.ModelProto, + verbose: bool = False, + ) -> onnx.ModelProto: + """Replace an exact generated 2x nearest-neighbor grid with Resize.""" + if _default_opset(model) < 11: + logger.warning( + "SurgeryPipe: gathernd-to-resize requires ONNX opset 11 or newer; skipping" + ) + return model + + index = _GraphIndex.build(model) + matches = [ + match + for node in model.graph.node + if node.op_type == "GatherND" + for match in [_match_resize_path(index, node)] + if match is not None + ] + if not matches: + return model + + allocator = _NameAllocator(model) + scales_name = allocator.new("surgery_resize_scales") + model.graph.initializer.append( + onnx.numpy_helper.from_array( + np.asarray([1.0, 1.0, 2.0, 2.0], dtype=np.float32), + scales_name, + ) + ) + + replacements: dict[int, onnx.NodeProto] = {} + remove_ids: set[int] = set() + prune_seeds: set[str] = set() + for match in matches: + resize = onnx.helper.make_node( + "Resize", + [match.source, "", scales_name], + [match.post_cast.output[0]], + name=allocator.new("surgery_nearest_resize"), + mode="nearest", + coordinate_transformation_mode="asymmetric", + nearest_mode="floor", + ) + replacements[id(match.pre_cast)] = resize + path_nodes = ( + match.pre_cast, + match.pre_transpose, + match.gather, + match.post_transpose, + match.post_cast, + ) + remove_ids.update(id(node) for node in path_nodes) + prune_seeds.update(name for node in path_nodes for name in node.input if name) + prune_seeds.add(match.indices) + if verbose: + logger.info( + " gathernd-to-resize: %s -> %s", + match.gather.name or match.gather.output[0], + resize.name, + ) + + rewritten: list[onnx.NodeProto] = [] + for node in model.graph.node: + replacement = replacements.get(id(node)) + if replacement is not None: + rewritten.append(replacement) + elif id(node) not in remove_ids: + rewritten.append(node) + del model.graph.node[:] + model.graph.node.extend(rewritten) + + _prune_dead_ancestors(model, prune_seeds) + _prune_stale_value_info(model) + logger.info( + "SurgeryPipe: gathernd-to-resize: replaced %d upsampling path(s)", + len(matches), + ) + return model + + # ----------------------------------------------------------------- + # silu-to-quick-gelu + # ----------------------------------------------------------------- + + def _silu_to_quick_gelu( + self, + model: onnx.ModelProto, + verbose: bool = False, + ) -> onnx.ModelProto: + """Replace exact x*Sigmoid(x) topology with QuickGelu(alpha=1).""" + index = _GraphIndex.build(model) + allocator = _NameAllocator(model) + replacements: dict[int, onnx.NodeProto] = {} + remove_ids: set[int] = set() + + for mul in list(model.graph.node): + if ( + mul.op_type != "Mul" + or not _is_default_domain(mul) + or len(mul.input) != 2 + or len(mul.output) != 1 + ): + continue + candidates: list[onnx.NodeProto] = [] + for sigmoid_input_index in (0, 1): + sigmoid = _producer(index, mul.input[sigmoid_input_index], "Sigmoid") + other_input = mul.input[1 - sigmoid_input_index] + sigmoid_output = _node_output(sigmoid) if sigmoid is not None else None + if ( + sigmoid is not None + and len(sigmoid.input) == 1 + and sigmoid.input[0] == other_input + and sigmoid_output not in index.graph_outputs + and _only_consumer(index, sigmoid_output, "Mul") is mul + ): + candidates.append(sigmoid) + if len(candidates) != 1: + continue + + sigmoid = candidates[0] + quick_gelu = onnx.helper.make_node( + "QuickGelu", + [sigmoid.input[0]], + [mul.output[0]], + name=allocator.new("surgery_silu_quick_gelu"), + domain="com.microsoft", + alpha=1.0, + ) + replacements[id(sigmoid)] = quick_gelu + remove_ids.update((id(sigmoid), id(mul))) + if verbose: + logger.info( + " silu-to-quick-gelu: %s + %s -> %s", + sigmoid.name or sigmoid.output[0], + mul.name or mul.output[0], + quick_gelu.name, + ) + + if not replacements: + return model + + rewritten: list[onnx.NodeProto] = [] + for node in model.graph.node: + replacement = replacements.get(id(node)) + if replacement is not None: + rewritten.append(replacement) + elif id(node) not in remove_ids: + rewritten.append(node) + del model.graph.node[:] + model.graph.node.extend(rewritten) + _ensure_ms_opset(model) + _prune_stale_value_info(model) + logger.info( + "SurgeryPipe: silu-to-quick-gelu: fused %d SiLU activation(s)", + len(replacements), + ) + return model + + # ----------------------------------------------------------------- + # scaled-matmul-to-fused-matmul + # ----------------------------------------------------------------- + + def _scaled_matmul_to_fused_matmul( + self, + model: onnx.ModelProto, + verbose: bool = False, + ) -> onnx.ModelProto: + """Replace MatMul followed by scalar Mul with FusedMatMul(alpha).""" + index = _GraphIndex.build(model) + allocator = _NameAllocator(model) + replacements: dict[int, onnx.NodeProto] = {} + remove_ids: set[int] = set() + prune_seeds: set[str] = set() + supported_types = { + int(onnx.TensorProto.FLOAT), + int(onnx.TensorProto.FLOAT16), + int(onnx.TensorProto.DOUBLE), + int(onnx.TensorProto.BFLOAT16), + } + + for matmul in list(model.graph.node): + if ( + matmul.op_type != "MatMul" + or not _is_default_domain(matmul) + or len(matmul.input) != 2 + or len(matmul.output) != 1 + or matmul.output[0] in index.graph_outputs + ): + continue + input_types = [index.element_types.get(name) for name in matmul.input] + if ( + input_types[0] is None + or input_types[0] != input_types[1] + or input_types[0] not in supported_types + ): + continue + + mul = _only_consumer(index, matmul.output[0], "Mul") + if mul is None or len(mul.input) != 2 or len(mul.output) != 1: + continue + if mul.input[0] == matmul.output[0]: + scale_name = mul.input[1] + elif mul.input[1] == matmul.output[0]: + scale_name = mul.input[0] + else: + continue + scale_values = _constant_array(index, scale_name) + scale_type = _constant_element_type(index, scale_name) + if scale_values is None or scale_values.size != 1 or scale_type != input_types[0]: + continue + try: + alpha = float(scale_values.reshape(-1)[0]) + except (TypeError, ValueError): + continue + if not np.isfinite(alpha) or float(np.float32(alpha)) != alpha: + continue + + fused_matmul = onnx.helper.make_node( + "FusedMatMul", + list(matmul.input), + [mul.output[0]], + name=allocator.new("surgery_scaled_fused_matmul"), + domain="com.microsoft", + alpha=alpha, + ) + replacements[id(matmul)] = fused_matmul + remove_ids.update((id(matmul), id(mul))) + prune_seeds.add(scale_name) + if verbose: + logger.info( + " scaled-matmul-to-fused-matmul: %s + %s -> %s (alpha=%s)", + matmul.name or matmul.output[0], + mul.name or mul.output[0], + fused_matmul.name, + alpha, + ) + + if not replacements: + return model + + rewritten: list[onnx.NodeProto] = [] + for node in model.graph.node: + replacement = replacements.get(id(node)) + if replacement is not None: + rewritten.append(replacement) + elif id(node) not in remove_ids: + rewritten.append(node) + del model.graph.node[:] + model.graph.node.extend(rewritten) + _ensure_ms_opset(model) + _prune_dead_ancestors(model, prune_seeds) + _prune_stale_value_info(model) + logger.info( + "SurgeryPipe: scaled-matmul-to-fused-matmul: fused %d MatMul(s)", + len(replacements), + ) + return model + + +class PreSurgeryPipe(SurgeryPipe): + """Apply proof-dependent graph surgeries before ORT rewrites the graph.""" + + name: ClassVar[str] = "pre-surgery" + capabilities: ClassVar[dict[str, Any]] = PRE_SURGERY_CAPABILITIES + + @classmethod + def build_config(cls, **kwargs: Any) -> SurgeryPipeConfig: + """Build a config containing only pre-optimization surgeries.""" + return SurgeryPipeConfig( + simplify_l2_normalization=kwargs.get("simplify_l2_normalization", False), + gathernd_to_resize=kwargs.get("gathernd_to_resize", False), + scaled_matmul_to_fused_matmul=kwargs.get("scaled_matmul_to_fused_matmul", False), + verbose=kwargs.get("verbose", False), + ) diff --git a/tests/unit/optim/pipes/test_pipe_surgery.py b/tests/unit/optim/pipes/test_pipe_surgery.py index 16df08d79..82abcefb3 100644 --- a/tests/unit/optim/pipes/test_pipe_surgery.py +++ b/tests/unit/optim/pipes/test_pipe_surgery.py @@ -343,6 +343,12 @@ def test_surgery_pipe_runs_last(self) -> None: # SurgeryPipe should be last in the list assert PIPES[-1].name == "surgery" + def test_proof_dependent_surgeries_run_first(self) -> None: + """Graph-pattern proofs must be checked before ORT rewrites the graph.""" + from winml.modelkit.optim.pipes import PIPES + + assert PIPES[0].name == "pre-surgery" + # ============================================================================= # UNTIE-CONSTANT-BATCHED-MATMUL TESTS diff --git a/tests/unit/optim/pipes/test_pipe_surgery_graph_reductions.py b/tests/unit/optim/pipes/test_pipe_surgery_graph_reductions.py new file mode 100644 index 000000000..fba91aa39 --- /dev/null +++ b/tests/unit/optim/pipes/test_pipe_surgery_graph_reductions.py @@ -0,0 +1,624 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Generated-model tests for exact and CUDA contrib graph surgeries.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import numpy as np +import onnx +import onnxruntime as ort +import pytest +from click.testing import CliRunner +from onnx import TensorProto, helper, numpy_helper +from onnx.reference import ReferenceEvaluator + +from winml.modelkit.commands.optimize import optimize +from winml.modelkit.optim import Optimizer +from winml.modelkit.optim.pipes import ( + SURGERY_CAPABILITIES, + SurgeryPipe, + SurgeryPipeConfig, +) + + +if TYPE_CHECKING: + from collections.abc import Callable + + +_NEW_CAPABILITIES = ( + "simplify-l2-normalization", + "gathernd-to-resize", + "silu-to-quick-gelu", + "scaled-matmul-to-fused-matmul", +) + + +def _run_ort(model: onnx.ModelProto, feeds: dict[str, np.ndarray]) -> list[np.ndarray]: + return ort.InferenceSession( + model.SerializeToString(), + providers=["CPUExecutionProvider"], + ).run(None, feeds) + + +def _attributes(node: onnx.NodeProto) -> dict[str, object]: + return {attribute.name: helper.get_attribute_value(attribute) for attribute in node.attribute} + + +def _make_l2_normalization_model( + *, + clip_min: float = 0.0, + axis: int = 1, + expand_batch: int = 1, +) -> onnx.ModelProto: + input_shape = [1, 3, 2, 2] + output_shape = [expand_batch, 3, 2, 2] + initializers = [ + numpy_helper.from_array(np.asarray([axis], dtype=np.int64), "axes"), + numpy_helper.from_array(np.asarray(clip_min, dtype=np.float32), "clip_min"), + numpy_helper.from_array(np.asarray(output_shape, dtype=np.int64), "expand_shape"), + ] + nodes = [ + helper.make_node( + "ReduceL2", + ["x", "axes"], + ["norm"], + name="reduce_l2", + keepdims=1, + ), + helper.make_node("Clip", ["norm", "clip_min"], ["clipped"], name="clip"), + helper.make_node( + "Expand", + ["clipped", "expand_shape"], + ["expanded"], + name="expand", + ), + helper.make_node("Div", ["x", "expanded"], ["y"], name="div"), + ] + graph = helper.make_graph( + nodes, + "l2_normalization", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, input_shape)], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, output_shape)], + initializer=initializers, + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 18)]) + return onnx.shape_inference.infer_shapes(model) + + +def _make_silu_model(*, scaled_sigmoid: bool = False) -> onnx.ModelProto: + nodes = [] + sigmoid_input = "x" + initializers = [] + if scaled_sigmoid: + initializers.append(numpy_helper.from_array(np.asarray(2.0, dtype=np.float32), "two")) + nodes.append(helper.make_node("Mul", ["x", "two"], ["scaled"], name="scale")) + sigmoid_input = "scaled" + nodes.extend( + [ + helper.make_node("Sigmoid", [sigmoid_input], ["sigmoid"], name="sigmoid"), + helper.make_node("Mul", ["x", "sigmoid"], ["y"], name="silu"), + ] + ) + graph = helper.make_graph( + nodes, + "silu", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, 3])], + initializer=initializers, + ) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 18)]) + + +def _make_scaled_matmul_model( + *, + scalar_scale: bool = True, + element_type: int = TensorProto.FLOAT, + scale_element_type: int | None = None, +) -> onnx.ModelProto: + scale_values = ( + np.asarray(0.125, dtype=np.float32) + if scalar_scale + else np.asarray([0.125, 0.25], dtype=np.float32) + ) + scale_element_type = scale_element_type or element_type + if scale_element_type == TensorProto.BFLOAT16: + scale = onnx.TensorProto(name="scale", data_type=TensorProto.BFLOAT16) + scale.dims.extend(scale_values.shape) + bits = np.right_shift(scale_values.view(np.uint32), 16).astype(np.uint16) + scale.raw_data = bits.tobytes() + else: + dtype = { + TensorProto.FLOAT16: np.float16, + TensorProto.FLOAT: np.float32, + TensorProto.DOUBLE: np.float64, + }[scale_element_type] + scale = numpy_helper.from_array(scale_values.astype(dtype), "scale") + graph = helper.make_graph( + [ + helper.make_node("MatMul", ["a", "b"], ["product"], name="matmul"), + helper.make_node("Mul", ["product", "scale"], ["y"], name="scale_product"), + ], + "scaled_matmul", + [ + helper.make_tensor_value_info("a", element_type, [2, 4]), + helper.make_tensor_value_info("b", element_type, [4, 2]), + ], + [helper.make_tensor_value_info("y", element_type, [2, 2])], + initializer=[scale], + ) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 18)]) + + +def _append_index_vector( + nodes: list[onnx.NodeProto], + dimension: str, + prefix: str, + *, + half_name: str, +) -> str: + doubled = f"{prefix}_doubled" + dimension_float = f"{prefix}_dimension_float" + doubled_float = f"{prefix}_doubled_float" + denominator = f"{prefix}_denominator" + ratio = f"{prefix}_ratio" + positions = f"{prefix}_positions" + offset_positions = f"{prefix}_offset_positions" + scaled_positions = f"{prefix}_scaled_positions" + indices = f"{prefix}_indices" + nodes.extend( + [ + helper.make_node("Mul", [dimension, "int_two"], [doubled], name=doubled), + helper.make_node( + "Cast", + [dimension], + [dimension_float], + name=dimension_float, + to=TensorProto.FLOAT, + ), + helper.make_node( + "Mul", + [dimension_float, "float_two"], + [denominator], + name=denominator, + ), + helper.make_node("Div", [dimension_float, denominator], [ratio], name=ratio), + helper.make_node( + "Cast", + [doubled], + [doubled_float], + name=doubled_float, + to=TensorProto.FLOAT, + ), + helper.make_node( + "Range", + ["float_zero", doubled_float, "float_one"], + [positions], + name=positions, + ), + helper.make_node( + "Add", + [positions, half_name], + [offset_positions], + name=offset_positions, + ), + helper.make_node( + "Mul", + [offset_positions, ratio], + [scaled_positions], + name=scaled_positions, + ), + helper.make_node( + "Cast", + [scaled_positions], + [indices], + name=indices, + to=TensorProto.INT64, + ), + ] + ) + return indices + + +def _make_gathernd_upsampling_model(*, half: float = 0.5) -> onnx.ModelProto: + input_shape = [1, 2, 2, 3] + output_shape = [1, 2, 4, 6] + initializers = [ + numpy_helper.from_array(np.asarray([1], dtype=np.int64), "batch_dim"), + numpy_helper.from_array(np.asarray([2], dtype=np.int64), "channel_dim"), + numpy_helper.from_array(np.asarray([-1], dtype=np.int64), "negative_one"), + numpy_helper.from_array(np.asarray(2, dtype=np.int64), "int_two"), + numpy_helper.from_array(np.asarray(0.0, dtype=np.float32), "float_zero"), + numpy_helper.from_array(np.asarray(1.0, dtype=np.float32), "float_one"), + numpy_helper.from_array(np.asarray(2.0, dtype=np.float32), "float_two"), + numpy_helper.from_array(np.asarray(half, dtype=np.float32), "half"), + ] + nodes = [ + helper.make_node("Shape", ["x"], ["height_vector"], name="height_shape", start=2, end=3), + helper.make_node("Squeeze", ["height_vector"], ["height"], name="height"), + helper.make_node("Shape", ["x"], ["width_vector"], name="width_shape", start=3, end=4), + helper.make_node("Squeeze", ["width_vector"], ["width"], name="width"), + helper.make_node( + "Concat", + ["batch_dim", "channel_dim", "height_vector", "width_vector"], + ["source_shape"], + name="source_shape", + axis=0, + ), + helper.make_node( + "Reshape", + ["x", "source_shape"], + ["source"], + name="source", + allowzero=1, + ), + ] + height_indices = _append_index_vector(nodes, "height", "height", half_name="half") + width_indices = _append_index_vector(nodes, "width", "width", half_name="half") + nodes.extend( + [ + helper.make_node( + "Unsqueeze", + [height_indices, "negative_one"], + ["height_column"], + name="height_column", + ), + helper.make_node( + "Max", + ["height_column", width_indices], + ["grid_extent"], + name="grid_extent", + ), + helper.make_node( + "Shape", + ["grid_extent"], + ["grid_shape"], + name="grid_shape", + start=0, + ), + helper.make_node( + "Expand", + ["height_column", "grid_shape"], + ["height_grid"], + name="height_grid", + ), + helper.make_node( + "Unsqueeze", + ["height_grid", "negative_one"], + ["height_component"], + name="height_component", + ), + helper.make_node( + "Expand", + [width_indices, "grid_shape"], + ["width_grid"], + name="width_grid", + ), + helper.make_node( + "Unsqueeze", + ["width_grid", "negative_one"], + ["width_component"], + name="width_component", + ), + helper.make_node( + "Concat", + ["height_component", "width_component"], + ["indices"], + name="indices", + axis=-1, + ), + helper.make_node( + "Cast", + ["source"], + ["source_float"], + name="source_float", + to=TensorProto.FLOAT, + ), + helper.make_node( + "Transpose", + ["source_float"], + ["spatial_first"], + name="spatial_first", + perm=[2, 3, 0, 1], + ), + helper.make_node( + "GatherND", + ["spatial_first", "indices"], + ["gathered"], + name="gathered", + batch_dims=0, + ), + helper.make_node( + "Transpose", + ["gathered"], + ["channels_first"], + name="channels_first", + perm=[2, 3, 0, 1], + ), + helper.make_node( + "Cast", + ["channels_first"], + ["y"], + name="output_cast", + to=TensorProto.FLOAT16, + ), + ] + ) + graph = helper.make_graph( + nodes, + "gathernd_upsampling", + [helper.make_tensor_value_info("x", TensorProto.FLOAT16, input_shape)], + [helper.make_tensor_value_info("y", TensorProto.FLOAT16, output_shape)], + initializer=initializers, + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 18)]) + return onnx.shape_inference.infer_shapes(model) + + +class TestGraphReductionCapabilities: + @pytest.mark.parametrize("name", _NEW_CAPABILITIES) + def test_capability_is_registered_and_disabled(self, name: str) -> None: + capability = SURGERY_CAPABILITIES[name] + assert capability.default is False + assert capability.ort_name is None + + @pytest.mark.parametrize("name", _NEW_CAPABILITIES) + def test_cli_lists_capability(self, name: str) -> None: + result = CliRunner().invoke(optimize, ["--list-capabilities"]) + assert result.exit_code == 0 + assert f"--enable-{name}" in result.output + + @pytest.mark.parametrize( + "kwarg", + [ + "simplify_l2_normalization", + "gathernd_to_resize", + "silu_to_quick_gelu", + "scaled_matmul_to_fused_matmul", + ], + ) + def test_each_capability_enables_surgery_pipe(self, kwarg: str) -> None: + config = SurgeryPipe.build_config(**{kwarg: True}) + assert getattr(config, kwarg) is True + assert SurgeryPipe.should_process(config) is True + + +class TestSimplifyL2Normalization: + def test_removes_only_clip_and_expand(self) -> None: + model = _make_l2_normalization_model() + result = SurgeryPipe().process( + model, + SurgeryPipeConfig(simplify_l2_normalization=True), + ) + + assert [node.op_type for node in result.graph.node] == ["ReduceL2", "Div"] + assert result.graph.node[-1].input == ["x", "norm"] + onnx.checker.check_model(result) + + @pytest.mark.parametrize( + "feeds", + [ + {"x": np.random.RandomState(1).randn(1, 3, 2, 2).astype(np.float32)}, + {"x": np.zeros((1, 3, 2, 2), dtype=np.float32)}, + ], + ) + def test_preserves_values_including_zero_norm(self, feeds: dict[str, np.ndarray]) -> None: + model = _make_l2_normalization_model() + result = SurgeryPipe().process( + model, + SurgeryPipeConfig(simplify_l2_normalization=True), + ) + expected = _run_ort(model, feeds)[0] + actual = _run_ort(result, feeds)[0] + np.testing.assert_allclose(actual, expected, rtol=0, atol=0, equal_nan=True) + + @pytest.mark.parametrize( + "builder", + [ + lambda: _make_l2_normalization_model(clip_min=1.0), + lambda: _make_l2_normalization_model(axis=2), + lambda: _make_l2_normalization_model(expand_batch=2), + ], + ) + def test_rejects_semantic_near_misses(self, builder: Callable[[], onnx.ModelProto]) -> None: + model = builder() + result = SurgeryPipe().process( + model, + SurgeryPipeConfig(simplify_l2_normalization=True), + ) + assert [node.op_type for node in result.graph.node] == [ + "ReduceL2", + "Clip", + "Expand", + "Div", + ] + + +class TestGatherNDToResize: + def test_replaces_exact_grid_and_prunes_grid_builders(self) -> None: + model = _make_gathernd_upsampling_model() + result = SurgeryPipe().process(model, SurgeryPipeConfig(gathernd_to_resize=True)) + + resize = next(node for node in result.graph.node if node.op_type == "Resize") + assert not any(node.op_type == "GatherND" for node in result.graph.node) + assert resize.input[0] == "source" + assert resize.output[0] == "y" + assert _attributes(resize) == { + "coordinate_transformation_mode": b"asymmetric", + "mode": b"nearest", + "nearest_mode": b"floor", + } + scales = next(init for init in result.graph.initializer if init.name == resize.input[2]) + np.testing.assert_array_equal( + numpy_helper.to_array(scales), + np.asarray([1.0, 1.0, 2.0, 2.0], dtype=np.float32), + ) + onnx.checker.check_model(result) + + def test_preserves_generated_fp16_values(self) -> None: + model = _make_gathernd_upsampling_model() + result = SurgeryPipe().process(model, SurgeryPipeConfig(gathernd_to_resize=True)) + feeds = {"x": np.random.RandomState(2).randn(1, 2, 2, 3).astype(np.float16)} + + expected = ReferenceEvaluator(model).run(None, feeds)[0] + actual = ReferenceEvaluator(result).run(None, feeds)[0] + np.testing.assert_array_equal(actual, expected) + + def test_rejects_different_coordinate_formula(self) -> None: + model = _make_gathernd_upsampling_model(half=0.25) + result = SurgeryPipe().process(model, SurgeryPipeConfig(gathernd_to_resize=True)) + assert any(node.op_type == "GatherND" for node in result.graph.node) + assert not any(node.op_type == "Resize" for node in result.graph.node) + + +class TestSiLUToQuickGelu: + def test_emits_quick_gelu_with_alpha_one(self) -> None: + model = _make_silu_model() + result = SurgeryPipe().process(model, SurgeryPipeConfig(silu_to_quick_gelu=True)) + + assert len(result.graph.node) == 1 + quick_gelu = result.graph.node[0] + assert (quick_gelu.domain, quick_gelu.op_type) == ("com.microsoft", "QuickGelu") + assert _attributes(quick_gelu)["alpha"] == pytest.approx(1.0) + assert [(opset.domain, opset.version) for opset in result.opset_import].count( + ("com.microsoft", 1) + ) == 1 + onnx.checker.check_model(result) + + def test_preserves_values(self) -> None: + model = _make_silu_model() + result = SurgeryPipe().process(model, SurgeryPipeConfig(silu_to_quick_gelu=True)) + feeds = {"x": np.random.RandomState(3).randn(2, 3).astype(np.float32)} + np.testing.assert_allclose( + _run_ort(result, feeds)[0], + _run_ort(model, feeds)[0], + rtol=1e-6, + atol=1e-7, + ) + + def test_rejects_scaled_sigmoid(self) -> None: + model = _make_silu_model(scaled_sigmoid=True) + result = SurgeryPipe().process(model, SurgeryPipeConfig(silu_to_quick_gelu=True)) + assert not any(node.op_type == "QuickGelu" for node in result.graph.node) + + +class TestScaledMatMulToFusedMatMul: + def test_emits_fused_matmul_and_removes_scale(self) -> None: + model = _make_scaled_matmul_model() + result = SurgeryPipe().process( + model, + SurgeryPipeConfig(scaled_matmul_to_fused_matmul=True), + ) + + assert len(result.graph.node) == 1 + fused = result.graph.node[0] + assert (fused.domain, fused.op_type) == ("com.microsoft", "FusedMatMul") + assert _attributes(fused)["alpha"] == pytest.approx(0.125) + assert all(initializer.name != "scale" for initializer in result.graph.initializer) + onnx.checker.check_model(result) + + def test_preserves_values(self) -> None: + model = _make_scaled_matmul_model() + result = SurgeryPipe().process( + model, + SurgeryPipeConfig(scaled_matmul_to_fused_matmul=True), + ) + feeds = { + "a": np.random.RandomState(4).randn(2, 4).astype(np.float32), + "b": np.random.RandomState(5).randn(4, 2).astype(np.float32), + } + np.testing.assert_array_equal(_run_ort(result, feeds)[0], _run_ort(model, feeds)[0]) + + def test_rejects_non_scalar_scale(self) -> None: + model = _make_scaled_matmul_model(scalar_scale=False) + result = SurgeryPipe().process( + model, + SurgeryPipeConfig(scaled_matmul_to_fused_matmul=True), + ) + assert [node.op_type for node in result.graph.node] == ["MatMul", "Mul"] + + def test_accepts_bfloat16_scalar(self) -> None: + model = _make_scaled_matmul_model(element_type=TensorProto.BFLOAT16) + result = SurgeryPipe().process( + model, + SurgeryPipeConfig(scaled_matmul_to_fused_matmul=True), + ) + assert [node.op_type for node in result.graph.node] == ["FusedMatMul"] + assert _attributes(result.graph.node[0])["alpha"] == pytest.approx(0.125) + + def test_rejects_scale_with_different_element_type(self) -> None: + model = _make_scaled_matmul_model( + element_type=TensorProto.FLOAT16, + scale_element_type=TensorProto.FLOAT, + ) + result = SurgeryPipe().process( + model, + SurgeryPipeConfig(scaled_matmul_to_fused_matmul=True), + ) + assert [node.op_type for node in result.graph.node] == ["MatMul", "Mul"] + + +def test_all_new_surgeries_are_idempotent() -> None: + models_and_configs = [ + (_make_l2_normalization_model(), SurgeryPipeConfig(simplify_l2_normalization=True)), + (_make_gathernd_upsampling_model(), SurgeryPipeConfig(gathernd_to_resize=True)), + (_make_silu_model(), SurgeryPipeConfig(silu_to_quick_gelu=True)), + ( + _make_scaled_matmul_model(), + SurgeryPipeConfig(scaled_matmul_to_fused_matmul=True), + ), + ] + for model, config in models_and_configs: + once = SurgeryPipe().process(model, config) + twice = SurgeryPipe().process(once, config) + assert once.SerializeToString() == twice.SerializeToString() + + +@pytest.mark.parametrize( + ("model_factory", "option", "removed_op", "added_op", "added_domain"), + [ + (_make_l2_normalization_model, "simplify_l2_normalization", "Clip", None, ""), + ( + _make_gathernd_upsampling_model, + "gathernd_to_resize", + "GatherND", + "Resize", + "", + ), + ( + _make_silu_model, + "silu_to_quick_gelu", + "Sigmoid", + "QuickGelu", + "com.microsoft", + ), + ( + _make_scaled_matmul_model, + "scaled_matmul_to_fused_matmul", + "MatMul", + "FusedMatMul", + "com.microsoft", + ), + ], +) +def test_surgery_survives_full_optimizer_pipeline( + model_factory: Callable[[], onnx.ModelProto], + option: str, + removed_op: str, + added_op: str | None, + added_domain: str, +) -> None: + result = Optimizer().optimize( + model_factory(), + constant_folding=False, + **{option: True}, + ) + + assert all(node.op_type != removed_op for node in result.graph.node) + if added_op is not None: + assert any( + node.op_type == added_op and node.domain == added_domain for node in result.graph.node + ) diff --git a/tests/unit/optim/test_optimizer.py b/tests/unit/optim/test_optimizer.py index 3095e30a7..fc5a65ca0 100644 --- a/tests/unit/optim/test_optimizer.py +++ b/tests/unit/optim/test_optimizer.py @@ -698,9 +698,9 @@ def test_resolve_dependencies_method(self) -> None: def test_registered_pipes_count(self) -> None: """Verify the expected number of pipes are registered.""" Optimizer._initialize_pipes() - # Currently: ORTGraphPipe, AlgebraicRewritePipe, RewritePipe, - # ORTFusionPipe, SurgeryPipe - assert len(Optimizer.pipes) == 5 + # Currently: PreSurgeryPipe, ORTGraphPipe, AlgebraicRewritePipe, + # RewritePipe, ORTFusionPipe, SurgeryPipe + assert len(Optimizer.pipes) == 6 def test_registered_pipe_names(self) -> None: """Verify expected pipe names are registered.""" @@ -710,4 +710,5 @@ def test_registered_pipe_names(self) -> None: assert "algebraic_rewrite" in names assert "ort_graph" in names assert "ort_fusion" in names + assert "pre-surgery" in names assert "surgery" in names diff --git a/vae_opt.md b/vae_opt.md new file mode 100644 index 000000000..626579683 --- /dev/null +++ b/vae_opt.md @@ -0,0 +1,484 @@ +# Wan 2.1 VAE Decoder ONNX Graph Optimization Report + +## Scope + +Model inspected directly: + +`C:\Users\hualxie\ComfyUI\models\vae_onnx\wan2.1_vae_decoder_fp16.onnx` + +External tensor data: + +`C:\Users\hualxie\ComfyUI\models\vae_onnx\wan2.1_vae_decoder_fp16.onnx.data` + +The model was analyzed without modifying either file. + +## Graph summary + +| Property | Value | +|---|---:| +| ONNX opset | 18 | +| Graph nodes | 9,558 | +| Initializers | 142 | +| Input | `z_tile`: FP16 `[1, 16, 21, h, w]` | +| Output | `sample_tile`: FP16 `[1, 3, 81, 8*h, 8*w]` | + +The largest operator groups are: + +| Operator | Count | +|---|---:| +| `Mul` | 2,359 | +| `Conv` | 927 | +| `Concat` | 845 | +| `Expand` | 794 | +| `Div` | 786 | +| `ReduceL2` | 780 | +| `Clip` | 780 | +| `Sigmoid` | 769 | +| `Add` | 381 | +| `Transpose` | 336 | +| `Reshape` | 279 | +| `Cast` | 160 | + +## Recommended exact rewrites + +### 1. Remove redundant `Clip` and `Expand` from channel L2 normalization + +**Occurrences:** 780 + +Every matching block has this form: + +```text +n = ReduceL2(x, axes=[1], keepdims=1) +n = Clip(n, min=0) +n = Expand(n, Shape(x)) +y = Div(x, n) +``` + +Replace the four nodes with: + +```text +n = ReduceL2(x, axes=[1], keepdims=1) +y = Div(x, n) +``` + +The scalar and learned-gamma multiplications after `Div` must remain in their +current order. + +#### Mathematical proof + +For a fixed batch and spatial location, let the channel vector be + +```text +x = (x_1, x_2, ..., x_C). +``` + +`ReduceL2` computes + +```text +n = sqrt(sum_j x_j^2) = ||x||_2. +``` + +Because a square is nonnegative, + +```text +x_j^2 >= 0 +``` + +and therefore + +```text +n >= 0. +``` + +Consequently, clipping the norm to a minimum of zero cannot change it: + +```text +max(n, 0) = n. +``` + +With `keepdims=1`, `n` already has a singleton channel dimension. `Expand` +only repeats that same norm across channels. ONNX `Div` supports +multidirectional broadcasting, so + +```text +x / Expand(n, Shape(x)) = x / n. +``` + +#### Reduction + +Each four-node block becomes two nodes: `780 * 2 = 1,560` fewer nodes. +Three shared shape-building `Concat` nodes also become dead, producing a total +reduction of **1,563 nodes**. + +#### Numerical caveat + +Keeping the original `ReduceL2` and `Div` kernels preserves their reduction, +rounding, and zero-norm behavior. In particular, an all-zero channel vector +still evaluates `0/0` instead of silently introducing an epsilon or returning +zero. + +The following multiplications must not be folded together if the current FP16 +rounding sequence must be retained: + +```text +(normalized_x * channel_scale) * gamma +``` + +Changing it to + +```text +normalized_x * (channel_scale * gamma) +``` + +is algebraically equal but can round differently in FP16. + +#### Conditional `LpNormalization` fusion + +For nonzero channel vectors, the two remaining nodes are algebraically equal +to: + +```text +y = LpNormalization(x, axis=1, p=2) +``` + +ONNX Runtime CUDA supports this operator for FP16 at opset 18. However, its +CUDA kernel intentionally returns zeros when the norm is zero, while the +original explicit `Div` produces NaNs from `0/0`. It also uses its own FP32 +reduction order for FP16 inputs. Therefore this one-node fusion is not globally +behavior-equivalent and must not be classified as exact unless nonzero norms +are proven or the zero-norm behavior change is explicitly accepted. + +### 2. Replace generated `GatherND` upsampling with nearest-neighbor `Resize` + +**Occurrences:** 73 + +Every spatial 2x upsampling block constructs integer indices and performs: + +```text +FP16 input + -> Cast(FP32) + -> Transpose + -> GatherND(generated indices) + -> Transpose + -> Cast(FP16) +``` + +The index-generation subgraphs are shared in three groups used by 11, 21, and +41 upsampling sites. + +Replace each data path with an ONNX `Resize` operating directly on FP16: + +```text +mode = "nearest" +coordinate_transformation_mode = "asymmetric" +nearest_mode = "floor" +axes = [2, 3] +scales = [2.0, 2.0] for the selected axes +``` + +#### Mathematical proof + +For output coordinate `i`, the existing graph computes + +```text +index(i) = int((i + 0.5) * (H / (2H))). +``` + +For nonzero `H`, + +```text +H / (2H) = 1/2, +``` + +so + +```text +index(i) = int((i + 0.5) / 2). +``` + +All generated coordinates are nonnegative, so conversion to `int64` is the +same as `floor`: + +```text +index(i) = floor((i + 0.5) / 2) = floor(i / 2). +``` + +Nearest-neighbor `Resize` with `asymmetric` coordinates and `floor` selection +uses + +```text +source(i) = floor(i / scale) = floor(i / 2). +``` + +Thus both formulations select exactly the same source element. The same proof +applies independently to height and width. + +The FP32 round trip is unnecessary because nearest-neighbor sampling performs +no arithmetic on tensor values; it only copies selected elements. Every FP16 +value is exactly representable in FP32, so FP16 -> FP32 -> FP16 does not alter +ordinary finite values. + +#### Reduction + +Replacing the 73 paths makes 433 old nodes dead, including the shared index +generation. Adding 73 `Resize` nodes gives a net reduction of **360 nodes**. + +#### Required constraints + +- Use `nearest`, not linear or cubic interpolation. +- Set `coordinate_transformation_mode="asymmetric"`. +- Set `nearest_mode="floor"`. +- Resize only spatial axes 2 and 3. +- Prefer constant `scales=[2.0, 2.0]`; CUDA expects the optional Resize + configuration inputs in CPU memory, and constant initializers avoid a + dynamic CPU shape subgraph. +- Preserve dynamic `h` and `w`; do not specialize them to one test shape. + +### Combined standard-ONNX result + +| Metric | Count | +|---|---:| +| Original nodes | 9,558 | +| Exact L2-normalization cleanup | 1,563 | +| Resize reduction | 360 | +| Resulting nodes | **7,635** | +| Total reduction | **1,923 (20.1%)** | + +Both recommended rewrites remain in the standard ONNX domain at opset 18. +If the conditional `LpNormalization` behavior is accepted, it removes another +780 nodes and produces the previously calculated 6,855-node graph. + +## Backend-specific fusions + +### 3. Fuse explicit SiLU + +**Occurrences:** 769 + +The graph expresses SiLU as: + +```text +s = Sigmoid(x) +y = Mul(x, s) +``` + +Mathematically, + +```text +sigmoid(x) = 1 / (1 + exp(-x)) +``` + +and therefore + +```text +y = x * sigmoid(x). +``` + +ONNX Runtime recognizes all 769 sites as +`com.microsoft::QuickGelu(alpha=1)`. Despite the operator name, with +`alpha=1` its expression is: + +```text +y = x * sigmoid(alpha*x) = x * sigmoid(x). +``` + +This reduces two graph nodes to one at each site. Applied after the two exact +standard rewrites, the graph would contain approximately **6,866 nodes**. +Applying the conditional `LpNormalization` fusion as well would produce the +previously calculated **6,086-node** graph. + +This is mathematically equivalent but uses a Microsoft-domain operator, so it +should only be emitted when the target execution provider supports it. + +### 4. Fuse attention score scaling into `MatMul` + +**Occurrences:** 11 + +The attention blocks contain: + +```text +s = MatMul(Q, K) +y = Mul(s, 0.051025390625) +``` + +A backend fused matrix multiplication with `alpha=0.051025390625` computes: + +```text +y = alpha * (Q @ K), +``` + +which is the same real-number expression. ONNX Runtime can represent this as +`com.microsoft::FusedMatMul`. + +This is backend-specific and may change floating-point rounding if scaling is +incorporated into the accumulation kernel. It should not be treated as a +bitwise-exact portable rewrite. + +## Opportunities requiring architectural transformation + +### Temporal `Concat -> Conv` unrolling + +There are 780 `Concat -> Conv` sites, but only 32 distinct convolution weight +tensors. Individual weights are reused 10, 20, or 40 times. This indicates +export-time temporal unrolling. + +Consolidating these repeated blocks could yield a much larger graph reduction, +but it is not a safe local fusion. The concatenations encode causal context, +zero initialization, and boundary handling. A correct rewrite must prove that +a single sequence-level convolution reproduces: + +- the same temporal receptive field; +- the same zero padding at the first frame; +- the same cached-state boundaries; +- the same output-frame ordering; and +- the same residual connections. + +This should be considered a separate optimization project rather than an +automatic peephole rewrite. + +## Rejected local rewrites + +- No adjacent inverse `Transpose -> Transpose` pairs were found. +- No adjacent redundant `Reshape -> Reshape` pairs were found. +- No adjacent redundant `Cast -> Cast` pairs were found. +- No exactly duplicated nodes were found. +- The 385 `Conv -> Add` paths are dynamic residual additions. Each convolution + already has a bias input, so these additions cannot be folded into Conv bias. +- Combining consecutive FP16 scalar/channel multiplications changes rounding + order and should not be called numerically exact. + +## Implemented `winml optimize` surgeries + +All four rewrites are implemented as independent, disabled-by-default flags: + +| Flag | Replacement | +|---|---| +| `--enable-simplify-l2-normalization` | Remove only the proven redundant `Clip` and `Expand` | +| `--enable-gathernd-to-resize` | Replace the exact generated 2x nearest-neighbor path with `Resize` | +| `--enable-silu-to-quick-gelu` | Replace `x * Sigmoid(x)` with `com.microsoft::QuickGelu(alpha=1)` | +| `--enable-scaled-matmul-to-fused-matmul` | Replace `MatMul -> Mul(scalar)` with `com.microsoft::FusedMatMul(alpha=scalar)` | + +The normalization, resize, and scaled-MatMul matchers run before ORT graph +optimization, while the SiLU fusion runs afterward so ORT does not inline +`QuickGelu` back into primitive operations. The matchers use graph topology, +types, shapes, constants, attributes, and consumer counts; they do not depend +on model, node, or tensor names. + +### Simple before-and-after examples + +#### 1. Simplify L2 normalization + +For a channel vector `x = [3, 4]`, `ReduceL2(x) = 5`. + +```text +Before: + norm = ReduceL2([3, 4]) = 5 + clipped = Clip(5, min=0) = 5 + expanded = Expand(5, shape=[2]) = [5, 5] + output = [3, 4] / [5, 5] = [0.6, 0.8] + +After: + norm = ReduceL2([3, 4]) = 5 + output = [3, 4] / 5 = [0.6, 0.8] +``` + +Broadcasting the scalar or singleton-channel norm gives the same result as +explicitly expanding it. + +#### 2. Replace `GatherND` upsampling with `Resize` + +For a 2x2 image, both graphs copy each value into a 2x2 output block: + +```text +Input: 2x nearest-neighbor output: + + [a b] [a a b b] + [c d] [a a b b] + [c c d d] + [c c d d] +``` + +The generated `GatherND` indices and `Resize(mode="nearest", +coordinate_transformation_mode="asymmetric", nearest_mode="floor")` both select +source coordinate `floor(output_coordinate / 2)`. + +#### 3. Fuse SiLU into `QuickGelu` + +For `x = 2`: + +```text +Before: + output = 2 * Sigmoid(2) + +After: + output = QuickGelu(2, alpha=1) + = 2 * Sigmoid(1 * 2) + = 2 * Sigmoid(2) +``` + +`QuickGelu` is exactly the same formula when `alpha=1`. + +#### 4. Fuse scaled `MatMul` + +```text +Before: + product = MatMul(A, B) + output = product * 0.5 + +After: + output = FusedMatMul(A, B, alpha=0.5) + = 0.5 * MatMul(A, B) +``` + +These are the same real-number expression. A fused kernel can change the final +floating-point rounding because it may apply `alpha` during accumulation. + +The target model was processed successfully with: + +```text +winml optimize -m C:\Users\hualxie\ComfyUI\models\vae_onnx\wan2.1_vae_decoder_fp16.onnx -o temp\wan2.1_vae_decoder_surgeries.onnx --overwrite --disable-constant-folding --enable-simplify-l2-normalization --enable-gathernd-to-resize --enable-silu-to-quick-gelu --enable-scaled-matmul-to-fused-matmul +``` + +The complete command pipeline produced 9,503 nodes. This differs from the +6,855-node direct-surgery count because the command's mandatory ORT graph pass +lowers and expands some standard operators. The requested transformations were +present in the saved graph: + +| Operator | Original | Optimized | +|---|---:|---:| +| `Clip` | 780 | 0 | +| `GatherND` | 73 | 0 | +| `Resize` | 0 | 73 | +| `Sigmoid` | 769 | 0 | +| `com.microsoft::QuickGelu` | 0 | 769 | +| `MatMul` | 22 | 11 | +| `com.microsoft::FusedMatMul` | 0 | 11 | + +The optimized model and its external-data sidecar pass ONNX full validation. + +## Validation requirements + +For each implemented rewrite: + +1. Run `onnx.checker.check_model` with the external tensor data available. +2. Compare original and optimized outputs with generated test inputs covering + multiple valid dynamic `h` and `w` values. +3. Include random inputs, all-zero inputs, large finite FP16 values, and values + near zero. +4. Compare each rewrite independently before testing the combined graph. +5. Test on the intended execution provider, not only CPU. +6. Require exact equality for the nearest-neighbor resize rewrite where the + provider preserves FP16 copies. +7. Use an explicitly approved FP16 tolerance for normalization and fused + kernels because mathematically equal reductions can accumulate in a + different order. + +## Recommended implementation order + +1. Implement `GatherND` upsampling to standard `Resize`; its value-selection + equivalence is strongest and it removes index-generation overhead. +2. Remove normalization `Clip` and `Expand`, retaining `ReduceL2 -> Div`. +3. Consider `LpNormalization` only if its zero-norm behavior is acceptable. +4. Enable SiLU and scaled-MatMul fusions only through provider capability + checks. +5. Evaluate temporal consolidation separately with full decoder equivalence + tests. From ec7937e3b755af6ac9734471be23a70a284f27d1 Mon Sep 17 00:00:00 2001 From: hualxie Date: Fri, 31 Jul 2026 10:51:27 +0800 Subject: [PATCH 4/4] Address optimize target review feedback --- src/winml/modelkit/commands/optimize.py | 14 +++++----- src/winml/modelkit/optim/analysis.py | 8 +++++- src/winml/modelkit/optim/pipes/graph.py | 20 +++++++++++++-- tests/unit/commands/test_optimize_cli.py | 16 ++++++++++++ tests/unit/optim/pipes/test_pipe_config.py | 30 ++++++++++++++++++---- tests/unit/optim/test_analysis.py | 4 ++- 6 files changed, 77 insertions(+), 15 deletions(-) diff --git a/src/winml/modelkit/commands/optimize.py b/src/winml/modelkit/commands/optimize.py index 8b1c01bd4..5ef06c1ed 100644 --- a/src/winml/modelkit/commands/optimize.py +++ b/src/winml/modelkit/commands/optimize.py @@ -596,7 +596,8 @@ def optimize( optimizer_kwargs["verbose"] = True if ep_device is not None: - assert target 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()}" @@ -626,11 +627,12 @@ def optimize( elif optimized_nodes > original_nodes: change_label = "increase" else: - change_label = "change" - node_info = ( - f"Nodes: {original_nodes} -> {optimized_nodes} " - f"({node_change:.1f}% {change_label})" - ) + 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: diff --git a/src/winml/modelkit/optim/analysis.py b/src/winml/modelkit/optim/analysis.py index f0b3fcb11..3d59adf98 100644 --- a/src/winml/modelkit/optim/analysis.py +++ b/src/winml/modelkit/optim/analysis.py @@ -346,8 +346,14 @@ def _iter_findings( if ep_device is not None and cap.ep_constraint is not None: from ..utils.constants import normalize_ep_name - target_ep = ep_device.device.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 diff --git a/src/winml/modelkit/optim/pipes/graph.py b/src/winml/modelkit/optim/pipes/graph.py index 6ea5feeae..4e5f4b3d1 100644 --- a/src/winml/modelkit/optim/pipes/graph.py +++ b/src/winml/modelkit/optim/pipes/graph.py @@ -597,11 +597,11 @@ def process(self, model: onnx.ModelProto, config: ORTGraphPipeConfig) -> onnx.Mo # Create session to trigger optimization try: if config.ep_device is None: - _ = ort.InferenceSession( + session = ort.InferenceSession( str(input_file), sess_opts, providers=["CPUExecutionProvider"] ) else: - _ = ort.InferenceSession(str(input_file), sess_opts) + session = ort.InferenceSession(str(input_file), sess_opts) except Exception as e: raise OptimizationError( f"ONNX Runtime optimization failed: {e}", @@ -613,6 +613,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) diff --git a/tests/unit/commands/test_optimize_cli.py b/tests/unit/commands/test_optimize_cli.py index feef461ef..6ad9d089b 100644 --- a/tests/unit/commands/test_optimize_cli.py +++ b/tests/unit/commands/test_optimize_cli.py @@ -220,6 +220,22 @@ def test_node_increase_reported(self, runner: CliRunner, tmp_path: Path) -> None assert result.exit_code == 0, result.output assert "20.0% increase" in result.output + def test_unchanged_node_count_reported(self, runner: CliRunner, tmp_path: Path) -> None: + model_file = tmp_path / "model.onnx" + model_file.touch() + model = _make_mock_model(num_nodes=10) + + with ( + patch(_LOAD_ONNX, return_value=model), + patch(_SAVE_ONNX), + patch(_OPTIMIZER) as mock_opt_cls, + ): + mock_opt_cls.return_value.optimize.return_value = model + result = runner.invoke(optimize, ["-m", str(model_file)]) + + assert result.exit_code == 0, result.output + assert "Nodes: 10 -> 10 (no change)" in result.output + def test_device_target_forwarded_to_optimizer( self, runner: CliRunner, tmp_path: Path ) -> None: diff --git a/tests/unit/optim/pipes/test_pipe_config.py b/tests/unit/optim/pipes/test_pipe_config.py index 3e46b6bc0..424200190 100644 --- a/tests/unit/optim/pipes/test_pipe_config.py +++ b/tests/unit/optim/pipes/test_pipe_config.py @@ -21,21 +21,19 @@ from __future__ import annotations import logging -from typing import TYPE_CHECKING from unittest.mock import MagicMock, patch +import pytest + from winml.modelkit.optim.pipes import ( GRAPH_CAPABILITIES, + OptimizationError, ORTGraphPipe, ORTGraphPipeConfig, PipeConfig, ) -if TYPE_CHECKING: - import pytest - - # ============================================================================= # TEST CONSTANTS - Capabilities used for testing (all default=False) # ============================================================================= @@ -164,12 +162,34 @@ def test_process_binds_resolved_ep_device(self) -> None: ), patch("winml.modelkit.session.lookup_device_spec", return_value=None), ): + inference_session.return_value.get_providers.return_value = [ + ep_device.device.ep_name + ] result = ORTGraphPipe().process(model, config) assert result is optimized_model session_options.add_provider_for_devices.assert_called_once_with([handle], {}) assert "providers" not in inference_session.call_args.kwargs + def test_process_rejects_provider_fallback(self) -> None: + model = MagicMock() + ep_device = MagicMock() + ep_device.device.ep_name = "DmlExecutionProvider" + session_options = MagicMock() + config = ORTGraphPipeConfig(ep_device=ep_device) + + with ( + patch("onnxruntime.SessionOptions", return_value=session_options), + patch("onnxruntime.InferenceSession") as inference_session, + patch("winml.modelkit.optim.pipes.graph.save_onnx"), + patch("winml.modelkit.session.lookup_device_spec", return_value=None), + pytest.raises(OptimizationError, match="was not activated"), + ): + inference_session.return_value.get_providers.return_value = [ + "CPUExecutionProvider" + ] + ORTGraphPipe().process(model, config) + def test_always_disabled_optimizers(self) -> None: """AttentionFusion and EmbedLayerNormFusion are ALWAYS disabled. diff --git a/tests/unit/optim/test_analysis.py b/tests/unit/optim/test_analysis.py index 1c7c59190..667ff7f09 100644 --- a/tests/unit/optim/test_analysis.py +++ b/tests/unit/optim/test_analysis.py @@ -407,7 +407,7 @@ def test_input_model_not_mutated(self) -> None: np.testing.assert_array_equal(before_value, after_value) def test_target_context_forwarded_and_ep_constraints_filtered( - self, monkeypatch + self, monkeypatch, caplog ) -> None: dml_cap = BoolCapability( name="dml-only", @@ -449,6 +449,7 @@ def process(model, config): ep_device = MagicMock() ep_device.device.ep_name = "DmlExecutionProvider" + caplog.set_level("DEBUG", logger="winml.modelkit.optim.analysis") monkeypatch.setattr("winml.modelkit.optim.pipes.PIPES", [TargetAwarePipe]) monkeypatch.setattr("winml.modelkit.onnx.infer_shapes", lambda model: model) @@ -461,6 +462,7 @@ def process(model, config): assert [finding.name for finding in findings] == ["dml-only"] assert captured_ep_devices assert all(value is ep_device for value in captured_ep_devices) + assert "Skipping capability 'cuda-only'" in caplog.text # =============================================================================