From 3bfb39dd6f7c4b6de1f8f2978d7170d8cfc1b917 Mon Sep 17 00:00:00 2001 From: fangyangci Date: Wed, 22 Jul 2026 21:22:25 +0800 Subject: [PATCH] squash: consolidate diff against origin/main --- .../cpu/audio-classification_fp16_config.json | 60 -- .../cpu/audio-classification_fp32_config.json | 41 - .../cpu/question-answering_fp16_config.json | 81 -- .../cpu/question-answering_fp32_config.json | 81 -- .../available_ops_all.json | 32 +- src/winml/modelkit/models/hf/__init__.py | 1 - src/winml/modelkit/models/hf/layoutlm.py | 68 -- src/winml/modelkit/onnx/io.py | 13 - .../modelkit/optim/capabilities/__init__.py | 2 - .../modelkit/optim/capabilities/algebraic.py | 32 - src/winml/modelkit/optim/pipes/__init__.py | 17 +- src/winml/modelkit/optim/pipes/algebraic.py | 930 ------------------ src/winml/modelkit/pattern/__init__.py | 2 + src/winml/modelkit/pattern/gelu_patterns.py | 8 + src/winml/modelkit/pattern/gemm_patterns.py | 20 +- .../modelkit/pattern/op_input_gen/__init__.py | 1 + .../op_input_gen/attention_input_generator.py | 59 ++ .../modelkit/pattern/transpose_patterns.py | 51 +- .../pattern/test_transpose_patterns.py | 49 + .../unit/export/test_onnx_config_overrides.py | 55 -- tests/unit/export/test_pytorch_export.py | 15 - tests/unit/optim/pipes/test_pipe_algebraic.py | 681 ------------- tests/unit/optim/test_optimizer.py | 6 +- 23 files changed, 213 insertions(+), 2092 deletions(-) delete mode 100644 examples/recipes/MIT_ast-finetuned-audioset-10-10-0.4593/cpu/cpu/audio-classification_fp16_config.json delete mode 100644 examples/recipes/MIT_ast-finetuned-audioset-10-10-0.4593/cpu/cpu/audio-classification_fp32_config.json delete mode 100644 examples/recipes/impira_layoutlm-document-qa/cpu/cpu/question-answering_fp16_config.json delete mode 100644 examples/recipes/impira_layoutlm-document-qa/cpu/cpu/question-answering_fp32_config.json delete mode 100644 src/winml/modelkit/models/hf/layoutlm.py delete mode 100644 src/winml/modelkit/optim/capabilities/algebraic.py delete mode 100644 src/winml/modelkit/optim/pipes/algebraic.py create mode 100644 src/winml/modelkit/pattern/op_input_gen/attention_input_generator.py delete mode 100644 tests/unit/optim/pipes/test_pipe_algebraic.py diff --git a/examples/recipes/MIT_ast-finetuned-audioset-10-10-0.4593/cpu/cpu/audio-classification_fp16_config.json b/examples/recipes/MIT_ast-finetuned-audioset-10-10-0.4593/cpu/cpu/audio-classification_fp16_config.json deleted file mode 100644 index 6ac10fb42..000000000 --- a/examples/recipes/MIT_ast-finetuned-audioset-10-10-0.4593/cpu/cpu/audio-classification_fp16_config.json +++ /dev/null @@ -1,60 +0,0 @@ -{ - "export": { - "opset_version": 17, - "batch_size": 1, - "export_params": true, - "do_constant_folding": true, - "verbose": false, - "dynamo": false, - "enable_hierarchy_tags": true, - "clean_onnx": false, - "hierarchy_tag_format": "full", - "input_tensors": [ - { - "name": "input_values", - "dtype": "float32", - "shape": [ - 1, - 1024, - 128 - ], - "value_range": [ - -1, - 1 - ] - } - ], - "output_tensors": [ - { - "name": "logits" - } - ] - }, - "optim": {}, - "quant": { - "mode": "fp16", - "samples": 10, - "calibration_method": "minmax", - "weight_type": "uint8", - "activation_type": "uint8", - "per_channel": false, - "symmetric": false, - "weight_symmetric": null, - "activation_symmetric": null, - "save_calibration": false, - "distribution": "uniform", - "seed": null, - "calibration_load_path": null, - "calibration_save_path": null, - "op_types_to_quantize": null, - "nodes_to_exclude": null, - "fp16_keep_io_types": true, - "fp16_op_block_list": null - }, - "compile": null, - "loader": { - "task": "audio-classification", - "model_class": "AutoModelForAudioClassification", - "model_type": "audio-spectrogram-transformer" - } -} diff --git a/examples/recipes/MIT_ast-finetuned-audioset-10-10-0.4593/cpu/cpu/audio-classification_fp32_config.json b/examples/recipes/MIT_ast-finetuned-audioset-10-10-0.4593/cpu/cpu/audio-classification_fp32_config.json deleted file mode 100644 index 439f9edee..000000000 --- a/examples/recipes/MIT_ast-finetuned-audioset-10-10-0.4593/cpu/cpu/audio-classification_fp32_config.json +++ /dev/null @@ -1,41 +0,0 @@ -{ - "export": { - "opset_version": 17, - "batch_size": 1, - "export_params": true, - "do_constant_folding": true, - "verbose": false, - "dynamo": false, - "enable_hierarchy_tags": true, - "clean_onnx": false, - "hierarchy_tag_format": "full", - "input_tensors": [ - { - "name": "input_values", - "dtype": "float32", - "shape": [ - 1, - 1024, - 128 - ], - "value_range": [ - -1, - 1 - ] - } - ], - "output_tensors": [ - { - "name": "logits" - } - ] - }, - "optim": {}, - "quant": null, - "compile": null, - "loader": { - "task": "audio-classification", - "model_class": "AutoModelForAudioClassification", - "model_type": "audio-spectrogram-transformer" - } -} diff --git a/examples/recipes/impira_layoutlm-document-qa/cpu/cpu/question-answering_fp16_config.json b/examples/recipes/impira_layoutlm-document-qa/cpu/cpu/question-answering_fp16_config.json deleted file mode 100644 index 0fe7d1a21..000000000 --- a/examples/recipes/impira_layoutlm-document-qa/cpu/cpu/question-answering_fp16_config.json +++ /dev/null @@ -1,81 +0,0 @@ -{ - "export": { - "opset_version": 17, - "batch_size": 1, - "export_params": true, - "do_constant_folding": true, - "verbose": false, - "dynamo": false, - "enable_hierarchy_tags": true, - "clean_onnx": false, - "hierarchy_tag_format": "full", - "input_tensors": [ - { - "name": "input_ids", - "dtype": "int32", - "shape": [ - 1, - 512 - ], - "value_range": [ - 0, - 50265 - ] - }, - { - "name": "bbox", - "dtype": "int32", - "shape": [ - 1, - 512, - 4 - ], - "value_range": [ - 0, - 1000 - ] - }, - { - "name": "attention_mask", - "dtype": "int32", - "shape": [ - 1, - 512 - ], - "value_range": [ - 0, - 2 - ] - }, - { - "name": "token_type_ids", - "dtype": "int32", - "shape": [ - 1, - 512 - ], - "value_range": [ - 0, - 1 - ] - } - ], - "output_tensors": [ - { - "name": "start_logits" - }, - { - "name": "end_logits" - } - ] - }, - "optim": { - "clamp_constant_values": true - }, - "quant": null, - "loader": { - "task": "question-answering", - "model_class": "LayoutLMForQuestionAnswering", - "model_type": "layoutlm" - } -} diff --git a/examples/recipes/impira_layoutlm-document-qa/cpu/cpu/question-answering_fp32_config.json b/examples/recipes/impira_layoutlm-document-qa/cpu/cpu/question-answering_fp32_config.json deleted file mode 100644 index 0fe7d1a21..000000000 --- a/examples/recipes/impira_layoutlm-document-qa/cpu/cpu/question-answering_fp32_config.json +++ /dev/null @@ -1,81 +0,0 @@ -{ - "export": { - "opset_version": 17, - "batch_size": 1, - "export_params": true, - "do_constant_folding": true, - "verbose": false, - "dynamo": false, - "enable_hierarchy_tags": true, - "clean_onnx": false, - "hierarchy_tag_format": "full", - "input_tensors": [ - { - "name": "input_ids", - "dtype": "int32", - "shape": [ - 1, - 512 - ], - "value_range": [ - 0, - 50265 - ] - }, - { - "name": "bbox", - "dtype": "int32", - "shape": [ - 1, - 512, - 4 - ], - "value_range": [ - 0, - 1000 - ] - }, - { - "name": "attention_mask", - "dtype": "int32", - "shape": [ - 1, - 512 - ], - "value_range": [ - 0, - 2 - ] - }, - { - "name": "token_type_ids", - "dtype": "int32", - "shape": [ - 1, - 512 - ], - "value_range": [ - 0, - 1 - ] - } - ], - "output_tensors": [ - { - "name": "start_logits" - }, - { - "name": "end_logits" - } - ] - }, - "optim": { - "clamp_constant_values": true - }, - "quant": null, - "loader": { - "task": "question-answering", - "model_class": "LayoutLMForQuestionAnswering", - "model_type": "layoutlm" - } -} diff --git a/src/winml/modelkit/analyze/utils/avalizble_ep_device_ops/available_ops_all.json b/src/winml/modelkit/analyze/utils/avalizble_ep_device_ops/available_ops_all.json index d40276597..2f135a038 100644 --- a/src/winml/modelkit/analyze/utils/avalizble_ep_device_ops/available_ops_all.json +++ b/src/winml/modelkit/analyze/utils/avalizble_ep_device_ops/available_ops_all.json @@ -2,7 +2,7 @@ "total": 118, "opset_range": [ 12, - 22 + 23 ], "ops_ai_onnx": { "Abs": { @@ -95,10 +95,10 @@ }, "Attention": { "version_key_map": { - "99999": "1n" + "23": "1n" }, - "start_version": 99999, - "valid_for_gene": false + "start_version": 23, + "valid_for_gene": true }, "AveragePool": { "version_key_map": { @@ -1023,8 +1023,28 @@ "Attention" ], "version_key_map": { - "99999": "8v" - } + "23": "8v" + }, + "check_for_eps": [ + "qnn_gpu", + "qnn_npu" + ] + }, + "SingleGeluPattern": { + "ops": [ + "Gelu" + ], + "version_key_map": { + "1": "aw" + }, + "check_for_eps": [ + "ov_cpu", + "ov_gpu", + "ov_npu", + "qnn_gpu", + "qnn_npu", + "nv_gpu" + ] }, "Conv2DInplaceLinear4DPattern": { "ops": [ diff --git a/src/winml/modelkit/models/hf/__init__.py b/src/winml/modelkit/models/hf/__init__.py index 0339b2b8e..3c5b0c12b 100644 --- a/src/winml/modelkit/models/hf/__init__.py +++ b/src/winml/modelkit/models/hf/__init__.py @@ -48,7 +48,6 @@ from .depth_anything import DepthAnythingIOConfig as _DepthAnythingIOConfig # triggers registration from .depth_pro import DepthProIOConfig as _DepthProIOConfig # triggers registration from .detr import DETR_CONFIG -from .layoutlm import LayoutLMQAIOConfig as _LayoutLMQAIOConfig # triggers registration from .marian import MARIAN_CONFIG from .marian import MODEL_CLASS_MAPPING as _MARIAN_CLASS_MAPPING from .marian import MarianDecoderIOConfig as _MarianDecoderIOConfig # triggers registration diff --git a/src/winml/modelkit/models/hf/layoutlm.py b/src/winml/modelkit/models/hf/layoutlm.py deleted file mode 100644 index 8085eb0ef..000000000 --- a/src/winml/modelkit/models/hf/layoutlm.py +++ /dev/null @@ -1,68 +0,0 @@ -# ------------------------------------------------------------------------- -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT License. -# -------------------------------------------------------------------------- -"""LayoutLM HuggingFace Model Configuration.""" - -from __future__ import annotations - -from typing import TYPE_CHECKING, Any, cast - -from optimum.exporters.onnx.model_configs import LayoutLMOnnxConfig -from optimum.utils import NormalizedTextConfig -from optimum.utils.input_generators import DummyBboxInputGenerator, DummyVisionInputGenerator - -from ...export import MaxLengthTextInputGenerator, register_onnx_overwrite - - -if TYPE_CHECKING: - import torch - - -class ZeroTokenTypeLayoutLMTextInputGenerator(MaxLengthTextInputGenerator): - """LayoutLM text dummy generator that keeps token_type_ids within type_vocab_size=1.""" - - def generate( - self, - input_name: str, - framework: str = "pt", - int_dtype: str = "int64", - float_dtype: str = "fp32", - ) -> torch.Tensor: - """Generate LayoutLM text inputs, replacing token_type_ids with zeros.""" - tensor = cast( - "torch.Tensor", - super().generate( - input_name, - framework=framework, - int_dtype=int_dtype, - float_dtype=float_dtype, - ), - ) - if input_name == "token_type_ids": - return tensor.new_zeros(tensor.shape) - return tensor - - -@register_onnx_overwrite("layoutlm", "question-answering", library_name="transformers") -class LayoutLMQAIOConfig(LayoutLMOnnxConfig): # type: ignore[misc] # optimum base is untyped - """LayoutLM question-answering OnnxConfig with bbox and safe token type IDs.""" - - # sequence_length is bound to the model's max_position_embeddings so - # MaxLengthTextInputGenerator emits full-length text inputs instead of - # Optimum's default of 16 (allow_new=True permits adding this mapping). - # We deliberately do NOT map max_2d_position_embeddings here: Optimum's - # DummyBboxInputGenerator hardcodes its coordinate range (its - # normalized_config.max_2d_position_embeddings read is commented out - # upstream), so such a mapping is inert and never becomes a sequence - # length. bbox coordinate bounds for the shipped recipe come from the - # recipe's `value_range` instead. - NORMALIZED_CONFIG_CLASS = NormalizedTextConfig.with_args( - sequence_length="max_position_embeddings", - allow_new=True, - ) - DUMMY_INPUT_GENERATOR_CLASSES: tuple[type[Any], ...] = ( - ZeroTokenTypeLayoutLMTextInputGenerator, - DummyVisionInputGenerator, - DummyBboxInputGenerator, - ) diff --git a/src/winml/modelkit/onnx/io.py b/src/winml/modelkit/onnx/io.py index d2ad85878..1c58362cf 100644 --- a/src/winml/modelkit/onnx/io.py +++ b/src/winml/modelkit/onnx/io.py @@ -106,19 +106,6 @@ def to_tensor(self) -> Any: lo, hi = self.value_range if torch_dtype.is_floating_point: return torch.rand(concrete_shape, dtype=torch_dtype) * (hi - lo) + lo - # bbox uses the HF convention: a [..., 4] tensor of box corners - # (x0, y0, x1, y1). We key on the input name because "bbox" is the - # canonical HF/Optimum name for 2D box coordinates; a model using a - # different name (e.g. "bboxes") simply falls through to the generic - # randint path below. Reorder so x0<=x1 and y0<=y1 to keep each box - # well-formed for layout models. - if self.name == "bbox" and len(concrete_shape) >= 1 and concrete_shape[-1] == 4: - coords = torch.randint(int(lo), int(hi), concrete_shape, dtype=torch_dtype) - x0 = torch.minimum(coords[..., 0], coords[..., 2]) - y0 = torch.minimum(coords[..., 1], coords[..., 3]) - x1 = torch.maximum(coords[..., 0], coords[..., 2]) - y1 = torch.maximum(coords[..., 1], coords[..., 3]) - return torch.stack((x0, y0, x1, y1), dim=-1) return torch.randint(int(lo), int(hi), concrete_shape, dtype=torch_dtype) # Fallback: no range info (backward compatible) diff --git a/src/winml/modelkit/optim/capabilities/__init__.py b/src/winml/modelkit/optim/capabilities/__init__.py index b219f1377..dcf954bf7 100644 --- a/src/winml/modelkit/optim/capabilities/__init__.py +++ b/src/winml/modelkit/optim/capabilities/__init__.py @@ -14,7 +14,6 @@ from . import ( activation, - algebraic, attention, conv, elimination, @@ -31,7 +30,6 @@ __all__ = [ "activation", - "algebraic", "attention", "conv", "elimination", diff --git a/src/winml/modelkit/optim/capabilities/algebraic.py b/src/winml/modelkit/optim/capabilities/algebraic.py deleted file mode 100644 index a6ac37655..000000000 --- a/src/winml/modelkit/optim/capabilities/algebraic.py +++ /dev/null @@ -1,32 +0,0 @@ -# ------------------------------------------------------------------------- -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT License. -# -------------------------------------------------------------------------- -"""Opt-in, exact algebraic graph-rewrite capabilities.""" - -from __future__ import annotations - -from ..registry import BoolCapability, CapabilityCategory - - -STATIC_SPLIT_TO_SLICE = BoolCapability( - name="static-split-to-slice", - ort_name=None, - description=( - "Replace statically bounded Split operations with standard Slice operations " - "while preserving output tensors" - ), - category=CapabilityCategory.REWRITE, - default=False, -) - -CONV_CHANNEL_AFFINE_FOLDING = BoolCapability( - name="conv-channel-affine-folding", - ort_name=None, - description=( - "Fold safe scalar or channel-wise constant Mul/Add branches into Conv " - "weights and bias; direct and static Split/Slice routing only" - ), - category=CapabilityCategory.REWRITE, - default=False, -) diff --git a/src/winml/modelkit/optim/pipes/__init__.py b/src/winml/modelkit/optim/pipes/__init__.py index 9315b71fa..6905a77b1 100644 --- a/src/winml/modelkit/optim/pipes/__init__.py +++ b/src/winml/modelkit/optim/pipes/__init__.py @@ -11,11 +11,6 @@ from typing import Any -from .algebraic import ( - ALGEBRAIC_CAPABILITIES, - AlgebraicRewritePipe, - AlgebraicRewritePipeConfig, -) from .base import BasePipe, OptimizationError, PipeConfig, caps_dict from .fusion import ORTFusionPipe, ORTFusionPipeConfig from .graph import GRAPH_CAPABILITIES, ORTGraphPipe, ORTGraphPipeConfig @@ -27,19 +22,12 @@ # - 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). -# - AlgebraicRewritePipe: Exact topology-based algebraic rewrites (after ORT folding). # - RewritePipe: Pattern-based subgraph rewriting (runs after ORT constant folding so that # shape constants are visible, but before ORTFusionPipe so normalised patterns are # available for transformer fusions). # - ORTFusionPipe: ORT transformer fusions (Python optimizer) # - SurgeryPipe: Post-optimization model surgery (runs last to clamp constants after folding) -PIPES: list[type[BasePipe]] = [ - ORTGraphPipe, - AlgebraicRewritePipe, - RewritePipe, - ORTFusionPipe, - SurgeryPipe, -] +PIPES: list[type[BasePipe]] = [ORTGraphPipe, RewritePipe, ORTFusionPipe, SurgeryPipe] def get_all_capabilities() -> dict[str, Any]: @@ -55,12 +43,9 @@ def get_all_capabilities() -> dict[str, Any]: __all__ = [ - "ALGEBRAIC_CAPABILITIES", "GRAPH_CAPABILITIES", "PIPES", "SURGERY_CAPABILITIES", - "AlgebraicRewritePipe", - "AlgebraicRewritePipeConfig", "BasePipe", "ORTFusionPipe", "ORTFusionPipeConfig", diff --git a/src/winml/modelkit/optim/pipes/algebraic.py b/src/winml/modelkit/optim/pipes/algebraic.py deleted file mode 100644 index 29aac986e..000000000 --- a/src/winml/modelkit/optim/pipes/algebraic.py +++ /dev/null @@ -1,930 +0,0 @@ -# ------------------------------------------------------------------------- -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT License. -# -------------------------------------------------------------------------- -"""Conservative, opt-in algebraic ONNX graph rewrites.""" - -from __future__ import annotations - -from dataclasses import dataclass -from typing import Any, ClassVar, cast - -import numpy as np -import onnx - -from ..capabilities import algebraic -from .base import BasePipe, PipeConfig, caps_dict - - -ALGEBRAIC_CAPABILITIES: dict[str, Any] = caps_dict( - algebraic.STATIC_SPLIT_TO_SLICE, - algebraic.CONV_CHANNEL_AFFINE_FOLDING, -) - - -@dataclass -class AlgebraicRewritePipeConfig(PipeConfig): - """Configuration for exact algebraic rewrites.""" - - static_split_to_slice: bool = False - conv_channel_affine_folding: bool = False - - -@dataclass -class _GraphIndex: - """Graph metadata required to identify statically bounded Split nodes.""" - - producers: dict[str, onnx.NodeProto] - consumers: dict[str, list[onnx.NodeProto]] - initializers: dict[str, onnx.TensorProto] - shapes: dict[str, tuple[int | 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 = {input_name for input_name in node.input if input_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 input_name in consumed_names: - consumers.setdefault(input_name, []).append(node) - - initializers = {initializer.name: initializer for initializer in graph.initializer} - shapes: dict[str, tuple[int | None, ...]] = {} - for value_info in (*graph.input, *graph.value_info, *graph.output): - shape = _value_info_shape(value_info) - if shape is not None: - shapes[value_info.name] = shape - for name, initializer in initializers.items(): - shapes.setdefault(name, tuple(int(dim) for dim in initializer.dims)) - - return cls( - producers=producers, - consumers=consumers, - initializers=initializers, - shapes=shapes, - graph_outputs={output.name for output in graph.output if output.name}, - ) - - -@dataclass -class _AffineCandidate: - """A safe affine branch associated with a Conv output channel interval.""" - - source_node: onnx.NodeProto - source_output_index: int - final_output: str - nodes: list[onnx.NodeProto] - start: int - end: int - scale: np.ndarray - offset: np.ndarray - - -class _NameAllocator: - """Allocate names without relying on optional or duplicated 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] - + [value.name for value in graph.value_info] - + [value.name for value in graph.output] - + [node.name for node in graph.node] - + [output for node in graph.node for output in 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 _value_info_shape(value_info: onnx.ValueInfoProto) -> tuple[int | None, ...] | None: - """Return a tensor shape, preserving unknown dimensions as ``None``.""" - if not value_info.type.HasField("tensor_type"): - return None - tensor_type = value_info.type.tensor_type - if not tensor_type.HasField("shape"): - return None - dimensions = [ - int(dimension.dim_value) if dimension.HasField("dim_value") else None - for dimension in tensor_type.shape.dim - ] - return tuple(dimensions) - - -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 _constant_array(index: _GraphIndex, name: str) -> np.ndarray | None: - """Read an initializer or a regular ONNX Constant value.""" - if not name: - return None - initializer = index.initializers.get(name) - if initializer is not None: - return np.asarray(onnx.numpy_helper.to_array(initializer)) - - producer = index.producers.get(name) - if producer is None or producer.op_type != "Constant": - return None - value = _attribute(producer, "value") - if value is not None: - try: - return np.asarray(onnx.numpy_helper.to_array(value)) - except (TypeError, ValueError): - return None - 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_ints(index: _GraphIndex, name: str) -> list[int] | None: - values = _constant_array(index, name) - if values is None or not np.issubdtype(values.dtype, np.integer): - return None - return [int(value) for value in values.reshape(-1).tolist()] - - -def _single_attribute_or_input_ints( - index: _GraphIndex, - node: onnx.NodeProto, - attribute_name: str, - input_index: int | None, -) -> tuple[list[int] | None, bool]: - """Read legacy attributes and newer constant inputs.""" - attribute_value = _attribute(node, attribute_name) - from_attribute = None - if attribute_value is not None: - try: - values = np.asarray(attribute_value) - if not np.issubdtype(values.dtype, np.integer): - return None, True - from_attribute = [int(value) for value in values.reshape(-1).tolist()] - except (TypeError, ValueError): - return None, True - - from_input = None - if input_index is not None and len(node.input) > input_index and node.input[input_index]: - from_input = _constant_ints(index, node.input[input_index]) - if from_input is None: - return None, True - - if from_attribute is not None and from_input is not None and from_attribute != from_input: - return None, True - return from_input if from_input is not None else from_attribute, False - - -def _node_output(node: onnx.NodeProto) -> str | None: - return node.output[0] if len(node.output) == 1 and node.output[0] else None - - -def _static_shape(index: _GraphIndex, name: str) -> tuple[int, ...] | None: - shape = index.shapes.get(name) - if shape is None or any(dimension is None for dimension in shape): - return None - return cast("tuple[int, ...]", shape) - - -def _new_initializer( - model: onnx.ModelProto, - allocator: _NameAllocator, - values: np.ndarray, - prefix: str, -) -> str: - name = allocator.new(prefix) - model.graph.initializer.append(onnx.numpy_helper.from_array(np.asarray(values), name)) - return name - - -def _remove_nodes(model: onnx.ModelProto, nodes: set[int]) -> None: - remaining = [node for node in model.graph.node if id(node) not in nodes] - del model.graph.node[:] - model.graph.node.extend(remaining) - - -def _captured_tensor_names(graph: onnx.GraphProto) -> set[str]: - """Return names a nested graph resolves from an enclosing scope.""" - 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 = {input_name for node in graph.node for input_name in node.input if input_name} - referenced.update(output.name for output in graph.output if output.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 or its nested subgraphs.""" - referenced = {value.name for value in (*graph.input, *graph.output) if value.name} - for node in graph.node: - referenced.update(input_name for input_name in node.input if input_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 _prune_unused_initializers(model: onnx.ModelProto) -> None: - used = _referenced_tensor_names(model.graph) - remaining = [initializer for initializer in model.graph.initializer if initializer.name in used] - del model.graph.initializer[:] - model.graph.initializer.extend(remaining) - - -def _prune_stale_value_info(model: onnx.ModelProto) -> None: - """Remove value metadata for tensors no longer defined by the graph.""" - graph = model.graph - defined = {value.name for value in graph.input 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 _prune_generated_slices(model: onnx.ModelProto, introduced: set[str]) -> None: - """Remove generated Slice nodes whose outputs are entirely dead.""" - while True: - index = _GraphIndex.build(model) - removable = { - id(node) - for node in model.graph.node - if node.name in introduced - and all( - output and output not in index.graph_outputs and not index.consumers.get(output) - for output in node.output - ) - } - if not removable: - return - _remove_nodes(model, removable) - - -def _prune_dead_constant_nodes(model: onnx.ModelProto) -> None: - while True: - index = _GraphIndex.build(model) - removable = { - id(node) - for node in model.graph.node - if node.op_type == "Constant" - and all( - output and output not in index.graph_outputs and not index.consumers.get(output) - for output in node.output - ) - } - if not removable: - return - _remove_nodes(model, removable) - - -def _split_boundaries( - index: _GraphIndex, - node: onnx.NodeProto, - input_name: str, -) -> tuple[int, list[tuple[int, int]]] | None: - """Return a static Split axis and output boundaries.""" - input_shape = index.shapes.get(input_name) - if input_shape is None or len(node.output) == 0: - return None - - axis_value = _attribute(node, "axis", 0) - if not isinstance(axis_value, (int, np.integer)): - return None - axis = int(axis_value) - if axis < -len(input_shape) or axis >= len(input_shape): - return None - axis %= len(input_shape) - axis_size = input_shape[axis] - if axis_size is None: - return None - - split_values, split_conflict = _single_attribute_or_input_ints(index, node, "split", 1) - if split_conflict: - return None - if split_values is None: - if axis_size <= 0 or axis_size % len(node.output) != 0: - return None - split_values = [axis_size // len(node.output)] * len(node.output) - if len(split_values) != len(node.output) or any(value <= 0 for value in split_values): - return None - if sum(split_values) != axis_size: - return None - - boundaries: list[tuple[int, int]] = [] - start = 0 - for size in split_values: - boundaries.append((start, start + size)) - start += size - return axis, boundaries - - -def _slice_channel_boundary( - index: _GraphIndex, - node: onnx.NodeProto, - input_name: str, - channel_axis: int, -) -> tuple[int, int] | None: - """Read a Slice that selects a contiguous, full non-channel region.""" - if len(node.input) < 2: - return None - input_shape = _static_shape(index, input_name) - if input_shape is None or channel_axis >= len(input_shape): - return None - starts = _constant_ints(index, node.input[1]) - ends = _constant_ints(index, node.input[2]) if len(node.input) > 2 else None - axes = _constant_ints(index, node.input[3]) if len(node.input) > 3 else None - steps = _constant_ints(index, node.input[4]) if len(node.input) > 4 else None - if starts is None or ends is None: - return None - if axes is None: - axes = list(range(len(starts))) - if steps is None: - steps = [1] * len(starts) - if not (len(starts) == len(ends) == len(axes) == len(steps)): - return None - normalized_axes: list[int] = [] - for axis in axes: - if axis < -len(input_shape) or axis >= len(input_shape): - return None - normalized_axes.append(axis % len(input_shape)) - if len(set(normalized_axes)) != len(normalized_axes) or any(step != 1 for step in steps): - return None - - def normalize_bound(value: int, axis_size: int, *, is_end: bool) -> int: - if value < 0: - value += axis_size - if is_end and value > axis_size: - return axis_size - return max(0, min(value, axis_size)) - - channel_boundary: tuple[int, int] | None = None - for start_value, end_value, axis in zip(starts, ends, normalized_axes, strict=True): - axis_size = input_shape[axis] - start = normalize_bound(start_value, axis_size, is_end=False) - end = normalize_bound(end_value, axis_size, is_end=True) - if axis == channel_axis: - if end <= start: - return None - channel_boundary = (start, end) - elif start != 0 or end != axis_size: - return None - return channel_boundary - - -def _channel_affine_values( - values: np.ndarray, - output_shape: tuple[int, ...], - channels: int, -) -> np.ndarray | None: - """Convert a scalar or a provably channel-only broadcast to ``[C]``.""" - if not np.issubdtype(values.dtype, np.floating): - return None - if values.size == 1: - return np.full(channels, values.reshape(-1)[0], dtype=values.dtype) - if values.ndim > len(output_shape): - return None - - padded = (1,) * (len(output_shape) - values.ndim) + tuple(values.shape) - for axis, dimension in enumerate(padded): - if dimension not in (1, output_shape[axis]): - return None - if axis != 1 and dimension != 1: - return None - if padded[1] != channels: - return None - return np.asarray(values).reshape(channels) - - -def _affine_operand( - index: _GraphIndex, - node: onnx.NodeProto, - data_name: str, - output_shape: tuple[int, ...], - channels: int, -) -> np.ndarray | None: - constant_inputs = [ - name - for name in node.input - if name and name != data_name and _constant_array(index, name) is not None - ] - if len(constant_inputs) != 1 or len(node.input) != 2: - return None - values = _constant_array(index, constant_inputs[0]) - return None if values is None else _channel_affine_values(values, output_shape, channels) - - -def _channel_preserving_view_output( - index: _GraphIndex, - node: onnx.NodeProto, - input_name: str, - channels: int, -) -> str | None: - """Return a shape-only view output that preserves N/C order.""" - output_name = _node_output(node) - if ( - output_name is None - or len(node.input) == 0 - or node.input[0] != input_name - or node.op_type not in {"Reshape", "Squeeze", "Unsqueeze"} - ): - return None - input_shape = _static_shape(index, input_name) - output_shape = _static_shape(index, output_name) - if ( - input_shape is None - or output_shape is None - or len(input_shape) < 2 - or len(output_shape) < 2 - or input_shape[:2] != output_shape[:2] - or input_shape[1] != channels - or np.prod(input_shape[2:], dtype=np.int64) != np.prod(output_shape[2:], dtype=np.int64) - ): - return None - - if node.op_type == "Reshape": - if ( - len(node.input) < 2 - or _constant_ints(index, node.input[1]) is None - or _attribute(node, "allowzero", 0) != 0 - ): - return None - else: - axes, conflict = _single_attribute_or_input_ints(index, node, "axes", 1) - if conflict or axes is None: - return None - return output_name - - -def _collect_affine_chain( - index: _GraphIndex, - first: onnx.NodeProto, - source_name: str, - output_shape: tuple[int, ...], - start: int, - end: int, - calculation_dtype: np.dtype[Any], -) -> _AffineCandidate | None: - """Collect a safe consecutive Mul/Add chain from one routed branch.""" - current = first - current_input = source_name - scale = np.ones(end - start, dtype=calculation_dtype) - offset = np.zeros(end - start, dtype=calculation_dtype) - matched: list[onnx.NodeProto] = [] - - while current.op_type in {"Mul", "Add"}: - if len(current.input) != 2 or current_input not in current.input: - return None - current_output = _node_output(current) - if current_output is None: - return None - values = _affine_operand(index, current, current_input, output_shape, end - start) - if values is None: - return None - values = values.astype(calculation_dtype, copy=False) - if current.op_type == "Mul": - scale *= values - offset *= values - else: - offset += values - matched.append(current) - - consumers = index.consumers.get(current_output, []) - if current_output in index.graph_outputs or len(consumers) != 1: - break - next_node = consumers[0] - if next_node.op_type not in {"Mul", "Add"}: - break - current_input = current_output - current = next_node - - final_output = _node_output(matched[-1]) if matched else None - if final_output is None or ( - final_output not in index.graph_outputs and len(index.consumers.get(final_output, [])) == 0 - ): - return None - return _AffineCandidate( - source_node=first, - source_output_index=0, - final_output=final_output, - nodes=matched, - start=start, - end=end, - scale=scale, - offset=offset, - ) - - -def _collect_routed_affine_candidates( - index: _GraphIndex, - source_node: onnx.NodeProto, - source_output_index: int, - start: int, - end: int, - calculation_dtype: np.dtype[Any], -) -> list[_AffineCandidate]: - """Collect affine leaves below safe views and disjoint channel slices.""" - if source_output_index >= len(source_node.output): - return [] - source_name = source_node.output[source_output_index] - if not source_name or source_name in index.graph_outputs: - return [] - - current_node = source_node - current_output_index = source_output_index - current_name = source_name - current_shape = _static_shape(index, current_name) - if current_shape is None or len(current_shape) < 2 or current_shape[1] != end - start: - return [] - - consumers = index.consumers.get(current_name, []) - while len(consumers) == 1: - view = consumers[0] - view_output = _channel_preserving_view_output( - index, - view, - current_name, - end - start, - ) - if view_output is None or current_name in index.graph_outputs: - break - current_node = view - current_output_index = 0 - current_name = view_output - current_shape = _static_shape(index, current_name) - if current_shape is None: - return [] - consumers = index.consumers.get(current_name, []) - - if current_name in index.graph_outputs: - return [] - if len(consumers) == 1 and consumers[0].op_type in {"Mul", "Add"}: - candidate = _collect_affine_chain( - index, - consumers[0], - current_name, - current_shape, - start, - end, - calculation_dtype, - ) - if candidate is None: - return [] - candidate.source_node = current_node - candidate.source_output_index = current_output_index - return [candidate] - - if not consumers or any(node.op_type != "Slice" for node in consumers): - return [] - routed_slices: list[tuple[onnx.NodeProto, int, int]] = [] - for routed_slice in consumers: - boundary = _slice_channel_boundary(index, routed_slice, current_name, 1) - if boundary is None: - return [] - routed_slices.append((routed_slice, *boundary)) - if any( - left_start < right_end and right_start < left_end - for position, (_, left_start, left_end) in enumerate(routed_slices) - for _, right_start, right_end in routed_slices[position + 1 :] - ): - return [] - - candidates: list[_AffineCandidate] = [] - for routed_slice, local_start, local_end in routed_slices: - candidates.extend( - _collect_routed_affine_candidates( - index, - routed_slice, - 0, - start + local_start, - start + local_end, - calculation_dtype, - ) - ) - return candidates - - -def _copy_conv_parameters( - model: onnx.ModelProto, - allocator: _NameAllocator, - conv: onnx.NodeProto, - scale: np.ndarray, - offset: np.ndarray, -) -> bool: - if len(conv.input) < 2: - return False - weight = next( - ( - initializer - for initializer in model.graph.initializer - if initializer.name == conv.input[1] - ), - None, - ) - if weight is None: - return False - weights = np.asarray(onnx.numpy_helper.to_array(weight)) - if weights.ndim < 1 or weights.shape[0] != len(scale): - return False - if not np.issubdtype(weights.dtype, np.floating): - return False - - if len(conv.input) > 2 and conv.input[2]: - bias = next( - ( - initializer - for initializer in model.graph.initializer - if initializer.name == conv.input[2] - ), - None, - ) - if bias is None: - return False - bias_values = np.asarray(onnx.numpy_helper.to_array(bias)) - if bias_values.ndim != 1 or len(bias_values) != len(scale): - return False - if not np.issubdtype(bias_values.dtype, np.floating): - return False - else: - bias_values = np.zeros(len(scale), dtype=weights.dtype) - - new_weights = weights * scale.reshape((len(scale),) + (1,) * (weights.ndim - 1)) - weight_name = _new_initializer( - model, - allocator, - np.asarray(new_weights, dtype=weights.dtype), - "algebraic_conv_weight", - ) - conv.input[1] = weight_name - new_bias = bias_values * scale + offset - bias_name = _new_initializer( - model, - allocator, - np.asarray(new_bias, dtype=bias_values.dtype), - "algebraic_conv_bias", - ) - if len(conv.input) > 2: - conv.input[2] = bias_name - else: - conv.input.append(bias_name) - return True - - -def _fold_channel_affine( - model: onnx.ModelProto, - allocator: _NameAllocator, -) -> None: - """Fold direct or static channel-routed affine branches after Conv.""" - index = _GraphIndex.build(model) - for original_conv in list(model.graph.node): - if ( - original_conv.op_type != "Conv" - or len(original_conv.output) != 1 - or not original_conv.output[0] - ): - continue - conv_output = original_conv.output[0] - conv = index.producers.get(conv_output) - if conv is None or conv.op_type != "Conv": - continue - conv_shape = _static_shape(index, conv_output) - if conv_shape is None or len(conv_shape) < 2: - continue - channels = conv_shape[1] - weight_initializer = index.initializers.get(conv.input[1]) if len(conv.input) > 1 else None - if weight_initializer is None: - continue - weight_dtype = onnx.numpy_helper.to_array(weight_initializer).dtype - if channels <= 0: - continue - calculation_dtype = np.result_type(weight_dtype, np.float32) - - route_name = conv_output - route_shape = conv_shape - route_source_node = conv - route_source_output_index = 0 - direct_consumers = index.consumers.get(route_name, []) - while len(direct_consumers) == 1: - view = direct_consumers[0] - view_output = _channel_preserving_view_output( - index, - view, - route_name, - channels, - ) - if view_output is None or route_name in index.graph_outputs: - break - route_name = view_output - next_route_shape = _static_shape(index, route_name) - if next_route_shape is None: - break - route_shape = next_route_shape - route_source_node = view - route_source_output_index = 0 - direct_consumers = index.consumers.get(route_name, []) - - candidates: list[_AffineCandidate] = [] - if route_name not in index.graph_outputs and len(direct_consumers) == 1: - direct = _collect_affine_chain( - index, - direct_consumers[0], - route_name, - route_shape, - 0, - channels, - calculation_dtype, - ) - if direct is not None: - direct.source_node = route_source_node - direct.source_output_index = route_source_output_index - candidates.append(direct) - - if not candidates and route_name not in index.graph_outputs and len(direct_consumers) == 1: - router = direct_consumers[0] - boundaries: list[tuple[int, int]] | None = None - if router.op_type == "Split": - split_info = _split_boundaries(index, router, route_name) - if split_info is not None and split_info[0] == 1: - boundaries = split_info[1] - elif router.op_type == "Slice": - boundary = _slice_channel_boundary(index, router, route_name, 1) - if boundary is not None: - boundaries = [boundary] - - if boundaries is not None and len(boundaries) == len(router.output): - for output_index, (start, end) in enumerate(boundaries): - candidates.extend( - _collect_routed_affine_candidates( - index, - router, - output_index, - start, - end, - calculation_dtype, - ) - ) - - if not candidates: - continue - if any( - left.start < right.end and right.start < left.end - for position, left in enumerate(candidates) - for right in candidates[position + 1 :] - ): - continue - if any( - output in index.graph_outputs - for candidate in candidates - for node in candidate.nodes[:-1] - for output in node.output - if output - ): - continue - - scale = np.ones(channels, dtype=calculation_dtype) - offset = np.zeros(channels, dtype=calculation_dtype) - for candidate in candidates: - scale[candidate.start : candidate.end] = candidate.scale - offset[candidate.start : candidate.end] = candidate.offset - if not _copy_conv_parameters(model, allocator, conv, scale, offset): - continue - - removed = {id(node) for candidate in candidates for node in candidate.nodes} - for candidate in candidates: - candidate.source_node.output[candidate.source_output_index] = candidate.final_output - _remove_nodes(model, removed) - index = _GraphIndex.build(model) - - -def _rewrite_static_splits( - model: onnx.ModelProto, - allocator: _NameAllocator, - introduced_nodes: set[str], -) -> None: - """Replace statically bounded Split nodes with input-form Slice nodes.""" - index = _GraphIndex.build(model) - opset = next( - (int(opset.version) for opset in model.opset_import if opset.domain in ("", "ai.onnx")), - 0, - ) - if opset and opset < 10: - return - - replacements: dict[int, list[onnx.NodeProto]] = {} - for split in list(model.graph.node): - if split.op_type != "Split" or len(split.input) < 1 or not split.input[0]: - continue - if any(not output for output in split.output): - continue - info = _split_boundaries(index, split, split.input[0]) - if info is None: - continue - axis, boundaries = info - replacement: list[onnx.NodeProto] = [] - for output_index, (start, end) in enumerate(boundaries): - starts_name = _new_initializer( - model, allocator, np.asarray([start], dtype=np.int64), "algebraic_slice_starts" - ) - ends_name = _new_initializer( - model, allocator, np.asarray([end], dtype=np.int64), "algebraic_slice_ends" - ) - axes_name = _new_initializer( - model, allocator, np.asarray([axis], dtype=np.int64), "algebraic_slice_axes" - ) - steps_name = _new_initializer( - model, allocator, np.asarray([1], dtype=np.int64), "algebraic_slice_steps" - ) - replacement_node = onnx.helper.make_node( - "Slice", - [split.input[0], starts_name, ends_name, axes_name, steps_name], - [split.output[output_index]], - name=allocator.new("algebraic_split_slice"), - ) - replacement.append(replacement_node) - introduced_nodes.add(replacement_node.name) - replacements[id(split)] = replacement - - if not replacements: - return - rewritten: list[onnx.NodeProto] = [] - for node in model.graph.node: - rewritten.extend(replacements.get(id(node), [node])) - del model.graph.node[:] - model.graph.node.extend(rewritten) - - -class AlgebraicRewritePipe(BasePipe[AlgebraicRewritePipeConfig]): - """Replace statically bounded Split nodes with Slice nodes.""" - - name: ClassVar[str] = "algebraic_rewrite" - capabilities: ClassVar[dict[str, Any]] = ALGEBRAIC_CAPABILITIES - - @classmethod - def build_config(cls, **kwargs: Any) -> AlgebraicRewritePipeConfig: - """Build the enabled algebraic rewrite configuration.""" - return AlgebraicRewritePipeConfig( - static_split_to_slice=kwargs.get("static_split_to_slice", False), - conv_channel_affine_folding=kwargs.get("conv_channel_affine_folding", False), - ) - - @classmethod - def should_process(cls, config: AlgebraicRewritePipeConfig) -> bool: - """Return whether any algebraic rewrite is enabled.""" - return config.static_split_to_slice or config.conv_channel_affine_folding - - def process( - self, - model: onnx.ModelProto, - config: AlgebraicRewritePipeConfig, - ) -> onnx.ModelProto: - """Rewrite eligible static Split nodes in a copy of the model.""" - if not self.should_process(config): - return model - - result = onnx.ModelProto() - result.CopyFrom(model) - allocator = _NameAllocator(result) - introduced_nodes: set[str] = set() - if config.conv_channel_affine_folding: - _fold_channel_affine(result, allocator) - if config.static_split_to_slice: - _rewrite_static_splits(result, allocator, introduced_nodes) - _prune_generated_slices(result, introduced_nodes) - _prune_dead_constant_nodes(result) - _prune_unused_initializers(result) - _prune_stale_value_info(result) - return result diff --git a/src/winml/modelkit/pattern/__init__.py b/src/winml/modelkit/pattern/__init__.py index 38f0112c4..94c672d25 100644 --- a/src/winml/modelkit/pattern/__init__.py +++ b/src/winml/modelkit/pattern/__init__.py @@ -49,6 +49,7 @@ Gelu4Pattern, Gelu4PatternInputGenerator, SingleGeluPattern, + SingleGeluPatternInputGenerator, ) from .gemm_patterns import ( MATMUL_ADD_SCHEMA, @@ -130,6 +131,7 @@ "ReshapeTransposeReshapeOverlyHighDimPattern", "ReshapeTransposeReshapeOverlyHighDimPatternInputGenerator", "SingleGeluPattern", + "SingleGeluPatternInputGenerator", "Skeleton", "SkeletonMatchResult", "SubgraphPattern", diff --git a/src/winml/modelkit/pattern/gelu_patterns.py b/src/winml/modelkit/pattern/gelu_patterns.py index 13c9586b4..0b2c6d8d1 100644 --- a/src/winml/modelkit/pattern/gelu_patterns.py +++ b/src/winml/modelkit/pattern/gelu_patterns.py @@ -444,3 +444,11 @@ class Gelu4PatternInputGenerator(PatternInputGenerator, get_runtime_checker_op(" pattern = Gelu4Pattern() registration_name = "Gelu4Pattern" + + +@register_pattern_input_generator +class SingleGeluPatternInputGenerator(PatternInputGenerator, get_runtime_checker_op("Gelu")): # type: ignore[misc] # dynamic base class (runtime-checker op) + """Input generator for native com.microsoft.Gelu single-op pattern.""" + + pattern = SingleGeluPattern() + registration_name = "SingleGeluPattern" diff --git a/src/winml/modelkit/pattern/gemm_patterns.py b/src/winml/modelkit/pattern/gemm_patterns.py index 5967c9535..0b089233e 100644 --- a/src/winml/modelkit/pattern/gemm_patterns.py +++ b/src/winml/modelkit/pattern/gemm_patterns.py @@ -19,6 +19,15 @@ from .op_input_gen import InputShapeConstraint, InputValueConstraint +def _is_scalar_constraint_value(value: Any) -> bool: + """Return True when the value represents a scalar literal constraint.""" + if value is None: + return False + if isinstance(value, np.ndarray): + return value.ndim == 0 + return np.isscalar(value) + + # Shared schema for MatMulAdd and ReshapeGemmReshape patterns # Both patterns perform the same computation: Y = MatMul(A, B) + C MATMUL_ADD_SCHEMA = PatternSchema( @@ -388,7 +397,7 @@ class GemmPatternInputGenerator(PatternInputGenerator): """PatternInputGenerator for Gemm patterns.""" # Typed as the base Pattern so subclasses can set a different concrete pattern. - pattern: Pattern = ReshapeGemmReshapePattern() + pattern = ReshapeGemmReshapePattern() def get_finite_attribute_sets(self) -> dict[str, list[Any]]: """Return finite attribute sets for ReshapeGemmReshape (none).""" @@ -447,9 +456,12 @@ def derive_properties(self, properties: dict) -> dict: else: item["C_dim"] = 0 - # Distinguish fixed-value bias constraints from shape-only constraints. - # This is finite (bool) and derives directly from existing row/query fields. - item["C_is_value_constraint"] = "C_value" in item and item["C_value"] is not None + # Distinguish scalar literal constraints from shape-based constraints. + # Vector/tensor constants should remain shape-constrained (not value-constrained). + c_value = item.get("C_value") + item["C_is_value_constraint"] = _is_scalar_constraint_value(c_value) and item[ + "C_dim" + ] == 0 return item def get_infinite_property_names(self) -> list[str]: diff --git a/src/winml/modelkit/pattern/op_input_gen/__init__.py b/src/winml/modelkit/pattern/op_input_gen/__init__.py index 0ae982453..94c718dcc 100644 --- a/src/winml/modelkit/pattern/op_input_gen/__init__.py +++ b/src/winml/modelkit/pattern/op_input_gen/__init__.py @@ -2,6 +2,7 @@ # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. # -------------------------------------------------------------------------- +from .attention_input_generator import AttentionInputGenerator from .binary_input_generator import * from .binary_like_input_generator import * from .constant_of_shape_input_generator import ConstantOfShapeInputGenerator diff --git a/src/winml/modelkit/pattern/op_input_gen/attention_input_generator.py b/src/winml/modelkit/pattern/op_input_gen/attention_input_generator.py new file mode 100644 index 000000000..be2ea0c81 --- /dev/null +++ b/src/winml/modelkit/pattern/op_input_gen/attention_input_generator.py @@ -0,0 +1,59 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Input generator for ai.onnx Attention operator (opset 23+). + +The generated combinations focus on the common 4D attention layout: +Q, K, V in BHSD format and optional attn_mask broadcastable to +(batch, heads, q_seq_len, kv_seq_len). +""" + +from .op_input_gen import InputShapeConstraint, OpInputGenerator, register_runtime_checker_op + + +@register_runtime_checker_op +class AttentionInputGenerator(OpInputGenerator): + """Input generator for Attention operator. + + Signature (opset 23): + Attention(Q, K, V, attn_mask?, past_key?, past_value?, ...) + + This generator currently targets 4D Q/K/V combinations that do not require + explicit q_num_heads/kv_num_heads attributes. + """ + + op_name = "Attention" + # Keep optional input expansion disabled so we can control optional-mask coverage + # explicitly and avoid invalid optional combinations. + expand_optionals = False + + def get_finite_attribute_sets(self) -> dict[str, list]: + """Attention attribute coverage is handled in input combinations.""" + return {} + + def get_input_and_infinite_attribute_combinations(self) -> list[dict[str, object]]: + """Return representative 4D input combinations for Attention.""" + return [ + # No mask: Q/K/V only. + { + "Q": InputShapeConstraint((1, 2, 4, 8)), + "K": InputShapeConstraint((1, 2, 6, 8)), + "V": InputShapeConstraint((1, 2, 6, 8)), + }, + # With mask (broadcast-exact shape for score tensor). + { + "Q": InputShapeConstraint((1, 2, 4, 8)), + "K": InputShapeConstraint((1, 2, 6, 8)), + "V": InputShapeConstraint((1, 2, 6, 8)), + "attn_mask": InputShapeConstraint((1, 2, 4, 6)), + }, + # Larger shape with explicit scale attribute. + { + "Q": InputShapeConstraint((2, 4, 8, 16)), + "K": InputShapeConstraint((2, 4, 8, 16)), + "V": InputShapeConstraint((2, 4, 8, 16)), + "attn_mask": InputShapeConstraint((2, 4, 8, 8)), + "scale": 1.0 / 16, + }, + ] \ No newline at end of file diff --git a/src/winml/modelkit/pattern/transpose_patterns.py b/src/winml/modelkit/pattern/transpose_patterns.py index b8ccd9cf5..9215aacae 100644 --- a/src/winml/modelkit/pattern/transpose_patterns.py +++ b/src/winml/modelkit/pattern/transpose_patterns.py @@ -338,7 +338,11 @@ def get_finite_attribute_sets(self) -> dict[str, list[Any]]: def get_input_and_infinite_attribute_combinations( self, ) -> list[dict[str, Any]]: - """Generate input and attribute combinations for testing.""" + """Generate representative input/attribute combinations for testing. + + The combinations intentionally cover both 3-D and 4-D output shapes, + which are commonly seen in real RTR subgraphs. + """ return [ { "data": InputShapeConstraint((1, 256, 256, 96)), @@ -346,12 +350,50 @@ def get_input_and_infinite_attribute_combinations( "perm": (0, 1, 3, 2, 4, 5), "output_shape": (1024, 8, 8, 96), }, + { + "data": InputShapeConstraint((1, 256, 256, 96)), + "transpose_shape": (1, 32, 8, 32, 8, 96), + "perm": (0, 1, 3, 2, 4, 5), + "output_shape": (1024, 64, 96), + }, { "data": InputShapeConstraint((1, 256, 256, 96)), "transpose_shape": (1, 32, 8, 32, 8, 96), "perm": (5, 4, 3, 2, 1, 0), "output_shape": (1024, 8, 8, 96), }, + { + "data": InputShapeConstraint((1, 256, 256, 96)), + "transpose_shape": (1, 32, 8, 32, 8, 96), + "perm": (5, 4, 3, 2, 1, 0), + "output_shape": (1, 65536, 96), + }, + { + "data": InputShapeConstraint((1, 224, 168, 128)), + "transpose_shape": (1, 32, 7, 24, 7, 128), + "perm": (0, 1, 3, 2, 4, 5), + "output_shape": (768, 49, 128), + }, + { + "data": InputShapeConstraint((1, 224, 168, 128)), + "transpose_shape": (1, 32, 24, 7, 7, 128), + "perm": (0, 1, 3, 2, 4, 5), + "output_shape": (-1, 224, 168, 128), + }, + { + "data": InputShapeConstraint((1, 192, 49, 256)), + "transpose_shape": (1, 16, 7, 12, 7, 256), + "perm": (0, 1, 3, 2, 4, 5), + "output_shape": (192, 49, 256), + }, + { + # PixelShuffle-style RTR variant observed in real models. + # This case provides merged_transpose_dim=5 coverage. + "data": InputShapeConstraint((1, 1, 3, 3, 4, 4)), + "transpose_shape": (1, 1, 3, 3, 4, 4), + "perm": (0, 1, 4, 2, 5, 3), + "output_shape": (1, 1, 12, 12), + }, ] def derive_properties(self, properties: dict) -> dict: @@ -367,7 +409,12 @@ def derive_properties(self, properties: dict) -> dict: def get_infinite_property_names(self) -> list[str]: """Return names of properties with infinite possible values.""" - return ["attr_transpose_shape", "attr_output_shape", "attr_perm", "data_shape"] + return [ + "attr_transpose_shape", + "attr_output_shape", + "attr_perm", + "data_shape", + ] @register_pattern_input_generator diff --git a/tests/unit/analyze/pattern/test_transpose_patterns.py b/tests/unit/analyze/pattern/test_transpose_patterns.py index b5aea72a0..aa5c0b694 100644 --- a/tests/unit/analyze/pattern/test_transpose_patterns.py +++ b/tests/unit/analyze/pattern/test_transpose_patterns.py @@ -20,7 +20,9 @@ ReshapeGemmReshapePattern, ReshapeTransposeReshapeLowDimPattern, ReshapeTransposeReshapeOverlyHighDimPattern, + get_pattern_input_generator, ) +from winml.modelkit.pattern.transpose_patterns import _compute_merged_transpose from .conftest import TEST_DOMAIN_VERSIONS @@ -244,3 +246,50 @@ def test_low_dim_pattern_id_differs_from_schema_name(self) -> None: pattern = ReshapeTransposeReshapeLowDimPattern() schema_based = f"SUBGRAPH/{pattern.get_schema().name}" assert pattern.pattern_id != schema_based + + +class TestRTRPatternInputGeneratorCoverage: + """Regression tests for RTR pattern input-generation coverage.""" + + def test_output_dim_coverage_includes_3d_and_4d(self) -> None: + """Generator should include both 3-D and 4-D output-shape variants.""" + generator_class = get_pattern_input_generator("ReshapeTransposeReshapeOverlyHighDimPattern") + gen = generator_class(domain_versions=TEST_DOMAIN_VERSIONS) + + combinations = gen.get_input_and_infinite_attribute_combinations() + output_dims: set[int] = set() + for item in combinations: + output_shape = item.get("output_shape") + assert isinstance(output_shape, tuple) + output_dims.add(len(output_shape)) + + assert 3 in output_dims + assert 4 in output_dims + + def test_output_dim_remains_finite_matching_property(self) -> None: + """output_dim should stay finite and participate in rule matching.""" + generator_class = get_pattern_input_generator("ReshapeTransposeReshapeOverlyHighDimPattern") + gen = generator_class(domain_versions=TEST_DOMAIN_VERSIONS) + + infinite_properties = gen.get_infinite_property_names() + assert "output_dim" not in infinite_properties + + def test_coverage_includes_merged_transpose_dim_5(self) -> None: + """Generator should include a case producing merged_transpose_dim=5.""" + generator_class = get_pattern_input_generator("ReshapeTransposeReshapeOverlyHighDimPattern") + gen = generator_class(domain_versions=TEST_DOMAIN_VERSIONS) + + combinations = gen.get_input_and_infinite_attribute_combinations() + merged_dims = set() + for item in combinations: + transpose_shape = item.get("transpose_shape") + perm = item.get("perm") + assert isinstance(transpose_shape, tuple) + assert isinstance(perm, tuple) + merged_shape, _ = _compute_merged_transpose( + transpose_shape, + perm, + ) + merged_dims.add(len(merged_shape)) + + assert 5 in merged_dims diff --git a/tests/unit/export/test_onnx_config_overrides.py b/tests/unit/export/test_onnx_config_overrides.py index 29478512c..f7ba15a39 100644 --- a/tests/unit/export/test_onnx_config_overrides.py +++ b/tests/unit/export/test_onnx_config_overrides.py @@ -57,7 +57,6 @@ class TestOnnxConfigRegistration: ("xlm-roberta", "fill-mask", "XLMRobertaIOConfig"), ("camembert", "fill-mask", "CamemBERTIOConfig"), ("mpnet", "fill-mask", "MPNetIOConfig"), - ("layoutlm", "question-answering", "LayoutLMQAIOConfig"), ("zoedepth", "depth-estimation", "ZoeDepthIOConfig"), ], ids=[ @@ -68,7 +67,6 @@ class TestOnnxConfigRegistration: "xlm-roberta", "camembert", "mpnet", - "layoutlm-qa", "zoedepth", ], ) @@ -133,59 +131,6 @@ def test_bert_io_specs_shape_matches(self, bert_config) -> None: ) -class TestLayoutLMQuestionAnsweringOverride: - """LayoutLM QA export must include bbox and safe token_type_ids.""" - - def test_layoutlm_qa_dummy_inputs_include_bbox_and_zero_token_types(self) -> None: - """Dummy inputs must keep bbox while forcing token_type_ids to zero.""" - from transformers import LayoutLMConfig - - layoutlm_config = LayoutLMConfig( - vocab_size=100, - hidden_size=64, - num_hidden_layers=2, - num_attention_heads=2, - intermediate_size=128, - max_position_embeddings=32, - max_2d_position_embeddings=1024, - type_vocab_size=1, - ) - - inputs = generate_dummy_inputs("layoutlm", "question-answering", layoutlm_config) - - assert set(inputs) == {"input_ids", "bbox", "attention_mask", "token_type_ids"} - assert inputs["input_ids"].shape == (1, layoutlm_config.max_position_embeddings) - assert inputs["bbox"].shape == (1, layoutlm_config.max_position_embeddings, 4) - assert inputs["token_type_ids"].shape == (1, layoutlm_config.max_position_embeddings) - assert inputs["token_type_ids"].max().item() == 0 - - def test_layoutlm_qa_io_specs_include_span_outputs(self) -> None: - """LayoutLM QA specs expose document bbox input and span logits outputs.""" - from transformers import LayoutLMConfig - - layoutlm_config = LayoutLMConfig( - vocab_size=100, - hidden_size=64, - num_hidden_layers=2, - num_attention_heads=2, - intermediate_size=128, - max_position_embeddings=32, - max_2d_position_embeddings=1024, - type_vocab_size=1, - ) - - specs = resolve_io_specs("layoutlm", "question-answering", layoutlm_config) - - assert specs["input_names"] == [ - "input_ids", - "bbox", - "attention_mask", - "token_type_ids", - ] - assert specs["input_shapes"] == [(1, 32), (1, 32, 4), (1, 32), (1, 32)] - assert specs["output_names"] == ["start_logits", "end_logits"] - - # ============================================================================= # Class 3: _adjust_position_embeddings unit tests # ============================================================================= diff --git a/tests/unit/export/test_pytorch_export.py b/tests/unit/export/test_pytorch_export.py index f22fc47b6..5b6aeae26 100644 --- a/tests/unit/export/test_pytorch_export.py +++ b/tests/unit/export/test_pytorch_export.py @@ -122,21 +122,6 @@ def test_int64_tensor(self) -> None: assert t.dtype == torch.int64 assert (t == 1).all() - def test_bbox_tensor_generates_ordered_coordinates_within_range(self) -> None: - spec = InputTensorSpec( - name="bbox", - dtype="int32", - shape=(2, 128, 4), - value_range=(0, 1000), - ) - t = spec.to_tensor() - assert t.dtype == torch.int32 - assert t.shape == (2, 128, 4) - assert int(t.min()) >= 0 - assert int(t.max()) < 1000 - assert torch.all(t[:, :, 2] >= t[:, :, 0]) - assert torch.all(t[:, :, 3] >= t[:, :, 1]) - def test_no_dtype_defaults_float(self) -> None: spec = InputTensorSpec(name="x", shape=(1, 10)) t = spec.to_tensor() diff --git a/tests/unit/optim/pipes/test_pipe_algebraic.py b/tests/unit/optim/pipes/test_pipe_algebraic.py deleted file mode 100644 index e9a977838..000000000 --- a/tests/unit/optim/pipes/test_pipe_algebraic.py +++ /dev/null @@ -1,681 +0,0 @@ -# ------------------------------------------------------------------------- -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT License. -# -------------------------------------------------------------------------- -"""Generated-graph tests for static Split-to-Slice rewriting.""" - -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 winml.modelkit.commands.optimize import optimize -from winml.modelkit.optim import get_all_capabilities, optimize_onnx -from winml.modelkit.optim.pipes import ( - PIPES, - AlgebraicRewritePipe, - AlgebraicRewritePipeConfig, -) - - -if TYPE_CHECKING: - from collections.abc import Sequence - - -def _tensor(name: str, values: np.ndarray) -> onnx.TensorProto: - return onnx.numpy_helper.from_array(np.asarray(values), name) - - -def _model( - nodes: Sequence[onnx.NodeProto], - inputs: Sequence[onnx.ValueInfoProto], - outputs: Sequence[onnx.ValueInfoProto], - initializers: Sequence[onnx.TensorProto], - value_info: Sequence[onnx.ValueInfoProto] = (), -) -> onnx.ModelProto: - graph = onnx.helper.make_graph( - list(nodes), - "generated_algebraic_graph", - list(inputs), - list(outputs), - initializer=list(initializers), - value_info=list(value_info), - ) - model = onnx.helper.make_model(graph, opset_imports=[onnx.helper.make_opsetid("", 17)]) - model.ir_version = 8 - return model - - -def _info(name: str, shape: Sequence[int | None]) -> onnx.ValueInfoProto: - return onnx.helper.make_tensor_value_info(name, onnx.TensorProto.FLOAT, list(shape)) - - -def _run(model: onnx.ModelProto, values: dict[str, np.ndarray]) -> list[np.ndarray]: - session = ort.InferenceSession( - model.SerializeToString(), - providers=["CPUExecutionProvider"], - ) - return session.run(None, values) - - -def _assert_valid_with_inferred_shapes(model: onnx.ModelProto) -> None: - onnx.checker.check_model(model) - inferred = onnx.shape_inference.infer_shapes(model) - assert len(inferred.graph.output) == len(model.graph.output) - - -class TestAlgebraicRegistration: - """Verify capability registration, flags, and pipe ordering.""" - - def test_capabilities_are_opt_in_and_independent(self) -> None: - capabilities = get_all_capabilities() - names = {"static-split-to-slice", "conv-channel-affine-folding"} - assert names <= capabilities.keys() - assert all(capabilities[name].default is False for name in names) - assert all( - capabilities[name].cli_flags() == (f"--enable-{name}", f"--disable-{name}") - for name in names - ) - - config = AlgebraicRewritePipe.build_config( - static_split_to_slice=True, - conv_channel_affine_folding=False, - ) - assert config.static_split_to_slice is True - assert config.conv_channel_affine_folding is False - - def test_cli_lists_algebraic_flag(self) -> None: - result = CliRunner().invoke(optimize, ["--list-capabilities"]) - assert result.exit_code == 0 - assert "--enable-static-split-to-slice" in result.output - assert "--enable-conv-channel-affine-folding" in result.output - - def test_pipe_is_after_ort_graph_and_before_cleanup(self) -> None: - names = [pipe.name for pipe in PIPES] - assert names.index("ort_graph") < names.index("algebraic_rewrite") - assert names.index("algebraic_rewrite") < names.index("surgery") - assert PIPES[names.index("algebraic_rewrite")] is AlgebraicRewritePipe - assert not AlgebraicRewritePipe.should_process(AlgebraicRewritePipeConfig()) - - -class TestStaticSplitToSlice: - """Test static Split replacement using generated data.""" - - def test_positive_equivalence_and_preserved_outputs(self) -> None: - rng = np.random.default_rng(10) - x = onnx.helper.make_tensor_value_info("x", onnx.TensorProto.FLOAT, [1, 6, 2]) - outputs = [_info("left", [1, 2, 2]), _info("right", [1, 4, 2])] - split = onnx.helper.make_node( - "Split", - ["x", "split_sizes"], - ["left", "right"], - name="", - axis=1, - ) - model = _model( - [split], - [x], - outputs, - [_tensor("split_sizes", np.asarray([2, 4], dtype=np.int64))], - ) - transformed = AlgebraicRewritePipe().process( - model, - AlgebraicRewritePipeConfig(static_split_to_slice=True), - ) - - assert [node.op_type for node in transformed.graph.node] == ["Slice", "Slice"] - assert [node.output[0] for node in transformed.graph.node] == ["left", "right"] - assert [output.name for output in transformed.graph.output] == ["left", "right"] - assert "split_sizes" not in { - initializer.name for initializer in transformed.graph.initializer - } - _assert_valid_with_inferred_shapes(transformed) - values = {"x": rng.normal(size=(1, 6, 2)).astype(np.float32)} - for original, rewritten in zip(_run(model, values), _run(transformed, values), strict=True): - np.testing.assert_allclose(original, rewritten, rtol=0, atol=0) - - def test_equal_split_and_name_collisions_are_safe(self) -> None: - x = onnx.helper.make_tensor_value_info("x", onnx.TensorProto.FLOAT, [1, 4, 2]) - keep = _info("keep", [1, 2, 2]) - split = onnx.helper.make_node( - "Split", - ["x"], - ["part_a", "part_b"], - name="", - axis=1, - ) - identity = onnx.helper.make_node( - "Identity", - ["part_b"], - ["keep"], - name="algebraic_split_slice", - ) - model = _model( - [split, identity], - [x], - [keep, _info("part_a", [1, 2, 2])], - [], - value_info=[_info("part_a", [1, 2, 2]), _info("part_b", [1, 2, 2])], - ) - transformed = AlgebraicRewritePipe().process( - model, - AlgebraicRewritePipeConfig(static_split_to_slice=True), - ) - generated = [node for node in transformed.graph.node if node.op_type == "Slice"] - assert len(generated) == 2 - assert len({node.name for node in transformed.graph.node}) == len(transformed.graph.node) - assert all(node.name for node in generated) - assert {node.output[0] for node in generated} == {"part_a", "part_b"} - _assert_valid_with_inferred_shapes(transformed) - - def test_dynamic_equal_split_and_malformed_split_are_unchanged(self) -> None: - x = onnx.helper.make_tensor_value_info("x", onnx.TensorProto.FLOAT, [1, None, 2]) - dynamic_equal = onnx.helper.make_node("Split", ["x"], ["a", "b"], axis=1) - malformed = onnx.helper.make_node( - "Split", - ["x", "bad_sizes"], - ["c", "d"], - axis=1, - ) - model = _model( - [dynamic_equal, malformed], - [x], - [_info("a", [1, None, 2]), _info("b", [1, None, 2])], - [_tensor("bad_sizes", np.asarray([1, 1], dtype=np.int64))], - ) - transformed = AlgebraicRewritePipe().process( - model, - AlgebraicRewritePipeConfig(static_split_to_slice=True), - ) - assert [node.op_type for node in transformed.graph.node] == ["Split", "Split"] - - def test_dead_generated_slice_and_constants_are_pruned(self) -> None: - x = onnx.helper.make_tensor_value_info("x", onnx.TensorProto.FLOAT, [1, 4, 2]) - model = _model( - [onnx.helper.make_node("Split", ["x"], ["left", "unused"], axis=1)], - [x], - [_info("left", [1, 2, 2])], - [], - ) - transformed = AlgebraicRewritePipe().process( - model, - AlgebraicRewritePipeConfig(static_split_to_slice=True), - ) - assert [node.op_type for node in transformed.graph.node] == ["Slice"] - assert transformed.graph.node[0].output[0] == "left" - assert len(transformed.graph.initializer) == 4 - - def test_nested_subgraph_captures_keep_generated_slices_live(self) -> None: - x = onnx.helper.make_tensor_value_info("x", onnx.TensorProto.FLOAT, [1, 4, 2]) - then_branch = onnx.helper.make_graph( - [onnx.helper.make_node("Identity", ["left"], ["then_output"])], - "then_branch", - [], - [_info("then_output", [1, 2, 2])], - ) - else_branch = onnx.helper.make_graph( - [onnx.helper.make_node("Identity", ["right"], ["else_output"])], - "else_branch", - [], - [_info("else_output", [1, 2, 2])], - ) - model = _model( - [ - onnx.helper.make_node("Split", ["x"], ["left", "right"], axis=1), - onnx.helper.make_node( - "If", - ["condition"], - ["y"], - then_branch=then_branch, - else_branch=else_branch, - ), - ], - [x], - [_info("y", [1, 2, 2])], - [_tensor("condition", np.asarray(True, dtype=np.bool_))], - value_info=[_info("left", [1, 2, 2]), _info("right", [1, 2, 2])], - ) - values = {"x": np.arange(8, dtype=np.float32).reshape(1, 4, 2)} - transformed = AlgebraicRewritePipe().process( - model, - AlgebraicRewritePipeConfig(static_split_to_slice=True), - ) - assert [node.op_type for node in transformed.graph.node] == [ - "Slice", - "Slice", - "If", - ] - _assert_valid_with_inferred_shapes(transformed) - np.testing.assert_array_equal(_run(model, values), _run(transformed, values)) - - def test_public_optimize_path_is_idempotent(self) -> None: - x = onnx.helper.make_tensor_value_info("x", onnx.TensorProto.FLOAT, [1, 4, 2]) - model = _model( - [onnx.helper.make_node("Split", ["x"], ["a", "b"], name="", axis=1)], - [x], - [_info("a", [1, 2, 2]), _info("b", [1, 2, 2])], - [], - ) - transformed = optimize_onnx(model, static_split_to_slice=True) - second = optimize_onnx(transformed, static_split_to_slice=True) - assert all(node.op_type != "Split" for node in transformed.graph.node) - assert [ - (node.op_type, tuple(node.input), tuple(node.output), node.name) - for node in transformed.graph.node - ] == [ - (node.op_type, tuple(node.input), tuple(node.output), node.name) - for node in second.graph.node - ] - _assert_valid_with_inferred_shapes(second) - - -class TestConvChannelAffineFolding: - """Test conservative direct and channel-routed affine folding.""" - - @pytest.fixture - def affine_model(self) -> tuple[onnx.ModelProto, dict[str, np.ndarray]]: - rng = np.random.default_rng(11) - x = onnx.helper.make_tensor_value_info("x", onnx.TensorProto.FLOAT, [1, 2, 3, 3]) - nodes = [ - onnx.helper.make_node("Conv", ["x", "weights"], ["conv_out"], name=""), - onnx.helper.make_node("Mul", ["conv_out", "scale"], ["mul_out"], name=""), - onnx.helper.make_node("Add", ["offset", "mul_out"], ["y"], name=""), - ] - weights = rng.normal(size=(3, 2, 1, 1)).astype(np.float32) - scale = rng.uniform(0.5, 1.5, size=(1, 3, 1, 1)).astype(np.float32) - offset = rng.normal(size=(1, 3, 1, 1)).astype(np.float32) - model = _model( - nodes, - [x], - [_info("y", [1, 3, 3, 3])], - [_tensor("weights", weights), _tensor("scale", scale), _tensor("offset", offset)], - value_info=[ - _info("conv_out", [1, 3, 3, 3]), - _info("mul_out", [1, 3, 3, 3]), - ], - ) - return model, {"x": rng.normal(size=(1, 2, 3, 3)).astype(np.float32)} - - def test_direct_affine_folding_is_exact_and_adds_optional_bias( - self, - affine_model: tuple[onnx.ModelProto, dict[str, np.ndarray]], - ) -> None: - model, values = affine_model - transformed = AlgebraicRewritePipe().process( - model, - AlgebraicRewritePipeConfig(conv_channel_affine_folding=True), - ) - assert [node.op_type for node in transformed.graph.node] == ["Conv"] - assert transformed.graph.node[0].output[0] == "y" - assert len(transformed.graph.node[0].input) == 3 - assert {initializer.name for initializer in transformed.graph.initializer} == set( - transformed.graph.node[0].input[1:] - ) - assert not transformed.graph.value_info - _assert_valid_with_inferred_shapes(transformed) - np.testing.assert_allclose( - _run(model, values), - _run(transformed, values), - rtol=2e-5, - atol=2e-5, - ) - - def test_float64_affine_values_preserve_weight_precision(self) -> None: - shape = [1, 1, 1, 1] - - def tensor_info(name: str) -> onnx.ValueInfoProto: - return onnx.helper.make_tensor_value_info(name, onnx.TensorProto.DOUBLE, shape) - - weight = np.ones(shape, dtype=np.float64) - scale = np.asarray(1.0 + 2**-30, dtype=np.float64) - model = _model( - [ - onnx.helper.make_node("Conv", ["x", "weight"], ["conv_out"]), - onnx.helper.make_node("Mul", ["conv_out", "scale"], ["y"]), - ], - [tensor_info("x")], - [tensor_info("y")], - [_tensor("weight", weight), _tensor("scale", scale)], - value_info=[tensor_info("conv_out")], - ) - - transformed = AlgebraicRewritePipe().process( - model, - AlgebraicRewritePipeConfig(conv_channel_affine_folding=True), - ) - folded_weight = next( - initializer - for initializer in transformed.graph.initializer - if initializer.name == transformed.graph.node[0].input[1] - ) - folded_values = onnx.numpy_helper.to_array(folded_weight) - assert folded_values.dtype == np.float64 - np.testing.assert_array_equal(folded_values, weight * scale) - - def test_shared_conv_output_is_ineligible( - self, - affine_model: tuple[onnx.ModelProto, dict[str, np.ndarray]], - ) -> None: - model, _ = affine_model - model.graph.node.append(onnx.helper.make_node("Identity", ["conv_out"], ["other"])) - model.graph.output.append(_info("other", [1, 3, 3, 3])) - transformed = AlgebraicRewritePipe().process( - model, - AlgebraicRewritePipeConfig(conv_channel_affine_folding=True), - ) - assert [node.op_type for node in transformed.graph.node] == [ - "Conv", - "Mul", - "Add", - "Identity", - ] - - def test_routed_view_graph_output_is_ineligible(self) -> None: - rng = np.random.default_rng(19) - model = _model( - [ - onnx.helper.make_node("Conv", ["x", "weight"], ["conv_out"]), - onnx.helper.make_node("Split", ["conv_out"], ["left", "right"], axis=1), - onnx.helper.make_node("Reshape", ["left", "view_shape"], ["left_view"]), - onnx.helper.make_node("Mul", ["left_view", "scale"], ["left_scaled"]), - onnx.helper.make_node("Identity", ["right"], ["right_out"]), - ], - [_info("x", [1, 1, 2, 2])], - [ - _info("left_view", [1, 1, 1, 2, 2]), - _info("left_scaled", [1, 1, 1, 2, 2]), - _info("right_out", [1, 1, 2, 2]), - ], - [ - _tensor("weight", rng.normal(size=(2, 1, 1, 1)).astype(np.float32)), - _tensor("view_shape", np.asarray([1, 1, 1, 2, 2], dtype=np.int64)), - _tensor("scale", np.asarray(1.5, dtype=np.float32)), - ], - value_info=[ - _info("conv_out", [1, 2, 2, 2]), - _info("left", [1, 1, 2, 2]), - _info("right", [1, 1, 2, 2]), - ], - ) - values = {"x": rng.normal(size=(1, 1, 2, 2)).astype(np.float32)} - transformed = AlgebraicRewritePipe().process( - model, - AlgebraicRewritePipeConfig(conv_channel_affine_folding=True), - ) - assert any(node.op_type == "Mul" for node in transformed.graph.node) - _assert_valid_with_inferred_shapes(transformed) - for original, rewritten in zip( - _run(model, values), - _run(transformed, values), - strict=True, - ): - np.testing.assert_array_equal(original, rewritten) - - def test_static_split_branches_fold_without_overlapping_ranges(self) -> None: - rng = np.random.default_rng(12) - x = onnx.helper.make_tensor_value_info("x", onnx.TensorProto.FLOAT, [1, 1, 2, 2]) - nodes = [ - onnx.helper.make_node("Conv", ["x", "weights"], ["conv_out"]), - onnx.helper.make_node( - "Split", - ["conv_out", "sizes"], - ["first", "second"], - axis=1, - ), - onnx.helper.make_node("Mul", ["first", "scale_first"], ["first_scaled"]), - onnx.helper.make_node("Add", ["first_scaled", "offset_first"], ["first_out"]), - onnx.helper.make_node("Add", ["second", "offset_second"], ["second_out"]), - ] - model = _model( - nodes, - [x], - [_info("first_out", [1, 2, 2, 2]), _info("second_out", [1, 2, 2, 2])], - [ - _tensor("weights", rng.normal(size=(4, 1, 1, 1)).astype(np.float32)), - _tensor("sizes", np.asarray([2, 2], dtype=np.int64)), - _tensor("scale_first", rng.uniform(size=(1, 2, 1, 1)).astype(np.float32)), - _tensor("offset_first", rng.normal(size=(1, 2, 1, 1)).astype(np.float32)), - _tensor("offset_second", rng.normal(size=(1, 2, 1, 1)).astype(np.float32)), - ], - value_info=[ - _info("conv_out", [1, 4, 2, 2]), - _info("first", [1, 2, 2, 2]), - _info("second", [1, 2, 2, 2]), - _info("first_scaled", [1, 2, 2, 2]), - ], - ) - values = {"x": rng.normal(size=(1, 1, 2, 2)).astype(np.float32)} - transformed = AlgebraicRewritePipe().process( - model, - AlgebraicRewritePipeConfig(conv_channel_affine_folding=True), - ) - assert [node.op_type for node in transformed.graph.node] == ["Conv", "Split"] - _assert_valid_with_inferred_shapes(transformed) - np.testing.assert_allclose( - _run(model, values), - _run(transformed, values), - rtol=2e-5, - atol=2e-5, - ) - - def test_channel_preserving_views_and_nested_slices_fold(self) -> None: - rng = np.random.default_rng(15) - x = onnx.helper.make_tensor_value_info("x", onnx.TensorProto.FLOAT, [1, 1, 2, 2]) - nodes = [ - onnx.helper.make_node("Conv", ["x", "weights"], ["conv_out"]), - onnx.helper.make_node("Reshape", ["conv_out", "view_shape"], ["viewed"]), - onnx.helper.make_node("Split", ["viewed", "split_sizes"], ["first", "second"], axis=1), - onnx.helper.make_node("Squeeze", ["first", "squeeze_axes"], ["first_view"]), - onnx.helper.make_node("Mul", ["first_view", "first_scale"], ["first_scaled"]), - onnx.helper.make_node("Relu", ["first_scaled"], ["first_out"]), - onnx.helper.make_node( - "Slice", - ["second", "slice_a_starts", "slice_a_ends"], - ["second_a"], - ), - onnx.helper.make_node( - "Slice", - [ - "second", - "slice_b_starts", - "slice_b_ends", - "slice_b_axes", - "slice_b_steps", - ], - ["second_b"], - ), - onnx.helper.make_node("Mul", ["second_a", "second_a_scale"], ["second_a_out"]), - onnx.helper.make_node("Add", ["second_b", "second_b_offset"], ["second_b_out"]), - ] - model = _model( - nodes, - [x], - [ - _info("first_out", [1, 2, 2, 2]), - _info("second_a_out", [1, 1, 1, 2, 2]), - _info("second_b_out", [1, 1, 1, 2, 2]), - ], - [ - _tensor("weights", rng.normal(size=(4, 1, 1, 1)).astype(np.float32)), - _tensor("view_shape", np.asarray([1, 4, 1, 2, 2], dtype=np.int64)), - _tensor("split_sizes", np.asarray([2, 2], dtype=np.int64)), - _tensor("squeeze_axes", np.asarray([2], dtype=np.int64)), - _tensor("slice_a_starts", np.asarray([0, 0, 0, 0, 0], dtype=np.int64)), - _tensor("slice_a_ends", np.asarray([1, 1, 1, 2, 2], dtype=np.int64)), - _tensor("slice_b_starts", np.asarray([0, 1, 0, 0, 0], dtype=np.int64)), - _tensor( - "slice_b_ends", - np.full(5, np.iinfo(np.int64).max, dtype=np.int64), - ), - _tensor("slice_b_axes", np.asarray([0, 1, 2, 3, 4], dtype=np.int64)), - _tensor("slice_b_steps", np.ones(5, dtype=np.int64)), - _tensor("first_scale", np.asarray(1.25, dtype=np.float32)), - _tensor("second_a_scale", np.asarray(0.75, dtype=np.float32)), - _tensor("second_b_offset", np.asarray(-0.5, dtype=np.float32)), - ], - value_info=[ - _info("conv_out", [1, 4, 2, 2]), - _info("viewed", [1, 4, 1, 2, 2]), - _info("first", [1, 2, 1, 2, 2]), - _info("second", [1, 2, 1, 2, 2]), - _info("first_view", [1, 2, 2, 2]), - _info("first_scaled", [1, 2, 2, 2]), - _info("second_a", [1, 1, 1, 2, 2]), - _info("second_b", [1, 1, 1, 2, 2]), - ], - ) - values = {"x": rng.normal(size=(1, 1, 2, 2)).astype(np.float32)} - config = AlgebraicRewritePipeConfig(conv_channel_affine_folding=True) - transformed = AlgebraicRewritePipe().process(model, config) - second = AlgebraicRewritePipe().process(transformed, config) - - assert not any(node.op_type in {"Mul", "Add"} for node in transformed.graph.node) - assert transformed.SerializeToString() == second.SerializeToString() - _assert_valid_with_inferred_shapes(transformed) - for original, rewritten in zip( - _run(model, values), - _run(transformed, values), - strict=True, - ): - np.testing.assert_allclose(original, rewritten, rtol=2e-5, atol=2e-5) - - public = optimize_onnx( - model, - conv_channel_affine_folding=True, - static_split_to_slice=True, - ) - second_public = optimize_onnx( - public, - conv_channel_affine_folding=True, - static_split_to_slice=True, - ) - assert not any(node.op_type in {"Mul", "Add", "Split"} for node in public.graph.node) - assert [ - (node.op_type, tuple(node.input), tuple(node.output), node.name) - for node in public.graph.node - ] == [ - (node.op_type, tuple(node.input), tuple(node.output), node.name) - for node in second_public.graph.node - ] - _assert_valid_with_inferred_shapes(second_public) - for original, rewritten in zip( - _run(model, values), - _run(second_public, values), - strict=True, - ): - np.testing.assert_allclose(original, rewritten, rtol=2e-5, atol=2e-5) - - def test_multiple_independent_affine_matches_fold(self) -> None: - rng = np.random.default_rng(16) - nodes = [ - onnx.helper.make_node("Conv", ["x1", "weight1"], ["conv1"]), - onnx.helper.make_node("Mul", ["conv1", "scale1"], ["y1"]), - onnx.helper.make_node("Conv", ["x2", "weight2"], ["conv2"]), - onnx.helper.make_node("Add", ["conv2", "offset2"], ["y2"]), - ] - model = _model( - nodes, - [_info("x1", [1, 1, 2, 2]), _info("x2", [1, 1, 2, 2])], - [_info("y1", [1, 1, 2, 2]), _info("y2", [1, 1, 2, 2])], - [ - _tensor("weight1", rng.normal(size=(1, 1, 1, 1)).astype(np.float32)), - _tensor("scale1", np.asarray(1.5, dtype=np.float32)), - _tensor("weight2", rng.normal(size=(1, 1, 1, 1)).astype(np.float32)), - _tensor("offset2", np.asarray(-0.25, dtype=np.float32)), - ], - value_info=[_info("conv1", [1, 1, 2, 2]), _info("conv2", [1, 1, 2, 2])], - ) - values = { - "x1": rng.normal(size=(1, 1, 2, 2)).astype(np.float32), - "x2": rng.normal(size=(1, 1, 2, 2)).astype(np.float32), - } - transformed = AlgebraicRewritePipe().process( - model, - AlgebraicRewritePipeConfig(conv_channel_affine_folding=True), - ) - assert [node.op_type for node in transformed.graph.node] == ["Conv", "Conv"] - for original, rewritten in zip( - _run(model, values), - _run(transformed, values), - strict=True, - ): - np.testing.assert_allclose(original, rewritten, rtol=2e-5, atol=2e-5) - - def test_nested_subgraph_captures_make_affine_fold_ineligible( - self, - affine_model: tuple[onnx.ModelProto, dict[str, np.ndarray]], - ) -> None: - model, values = affine_model - branch_shape = [1, 3, 3, 3] - then_branch = onnx.helper.make_graph( - [onnx.helper.make_node("Identity", ["mul_out"], ["then_output"])], - "then_branch", - [], - [_info("then_output", branch_shape)], - ) - else_branch = onnx.helper.make_graph( - [onnx.helper.make_node("Identity", ["conv_out"], ["else_output"])], - "else_branch", - [], - [_info("else_output", branch_shape)], - ) - model.graph.initializer.append(_tensor("condition", np.asarray(True, dtype=np.bool_))) - model.graph.node.append( - onnx.helper.make_node( - "If", - ["condition"], - ["captured"], - then_branch=then_branch, - else_branch=else_branch, - ) - ) - model.graph.output.append(_info("captured", branch_shape)) - - transformed = AlgebraicRewritePipe().process( - model, - AlgebraicRewritePipeConfig(conv_channel_affine_folding=True), - ) - assert [node.op_type for node in transformed.graph.node] == [ - "Conv", - "Mul", - "Add", - "If", - ] - _assert_valid_with_inferred_shapes(transformed) - for original, rewritten in zip( - _run(model, values), - _run(transformed, values), - strict=True, - ): - np.testing.assert_allclose(original, rewritten, rtol=0, atol=0) - - def test_constant_attribute_affine_is_folded_and_pruned(self) -> None: - rng = np.random.default_rng(18) - model = _model( - [ - onnx.helper.make_node("Conv", ["x", "weight"], ["conv_out"]), - onnx.helper.make_node("Constant", [], ["scale"], value_float=1.5), - onnx.helper.make_node("Mul", ["conv_out", "scale"], ["y"]), - ], - [_info("x", [1, 1, 2, 2])], - [_info("y", [1, 1, 2, 2])], - [_tensor("weight", rng.normal(size=(1, 1, 1, 1)).astype(np.float32))], - value_info=[_info("conv_out", [1, 1, 2, 2])], - ) - values = {"x": rng.normal(size=(1, 1, 2, 2)).astype(np.float32)} - transformed = AlgebraicRewritePipe().process( - model, - AlgebraicRewritePipeConfig(conv_channel_affine_folding=True), - ) - assert [node.op_type for node in transformed.graph.node] == ["Conv"] - np.testing.assert_allclose( - _run(model, values), - _run(transformed, values), - rtol=2e-5, - atol=2e-5, - ) diff --git a/tests/unit/optim/test_optimizer.py b/tests/unit/optim/test_optimizer.py index 3095e30a7..1e51af2b7 100644 --- a/tests/unit/optim/test_optimizer.py +++ b/tests/unit/optim/test_optimizer.py @@ -698,16 +698,14 @@ 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: ORTGraphPipe, RewritePipe, ORTFusionPipe, SurgeryPipe + assert len(Optimizer.pipes) == 4 def test_registered_pipe_names(self) -> None: """Verify expected pipe names are registered.""" Optimizer._initialize_pipes() names = {pipe_class.name for pipe_class in Optimizer.pipes} assert "rewrite" in names - assert "algebraic_rewrite" in names assert "ort_graph" in names assert "ort_fusion" in names assert "surgery" in names